#!/usr/bin/env lua 
-- <!-- vim : set tabstop=2 shiftwidth=2 expandtab : -->
local the,help = {},[[

kah.lua : how to change your mind (using TPE + Bayes classifier)
(c) 2024 Tim Menzies (timm@ieee.org). BSD-2 license

USAGE: 
  chmod +x kah.lua
  ./kah.lua [OPTIONS]

OPTIONS:

  -B Bootstraps num.of bootstraps        = 256
  -c conf       statistical confidence   = .05
  -C Cliffs     threshjold for cliffs    = .195
  -h help       show help                = false
  -k k          bayes control            = 1
  -m m          bayes control            = 2
  -p p          distance coeffecient     = 2
  -r rseed      random seed              = 1234567891
  -s start      init number of samples   = 4
  -S Stop       max number of labellings = 30
  -H Holdout    testing hold our ratio   = .33
  -t train      data                     = ../../moot/optimize/misc/auto93.csv 
]]

local fmt, pop           = string.format, table.remove
local abs, cos, exp, log = math.abs, math.cos, math.exp, math.log
local max, min, pi, sqrt = math.max, math.min, math.pi, math.sqrt
local R                  = math.random

local SYM,NUM,DATA,COLS = {},{},{},{}
local big = 1E32
local fail, pass= "❌","✅"
      
----------------- ----------------- ----------------- ----------------- ------------------
-- ## Utilities
-- ###  Lists

-- _push(list,atom) --> atom_
local function push(t,x) t[1+#t] = x; return x end

-- `split(list, n:int=#list) -> list,list`
local function split(t,n)
  local u,v = {},{}
  for j,x in pairs(t) do push(j <= n and u or v,x) end
  return u,v end

-- ### Random 

-- _any(list) --> Any_
local function any(t)    
  return t[R(#t)] end

-- _many(list, n=#list) --> list_
local function many(t,n,  u) 
  u={}; for i=1,(n or #t) do u[i] = any(t) end; return u end

-- _normal(mu:num=0, sd:num) --> num_
local function normal(mu,sd) 
   return (mu or 0)+(sd or 1)*sqrt(-2*log(R()))*cos(2*pi*R()) end
-- ### Sorting

-- _lt(atom) --> function_
local function lt(x) 
  return function(a,b) return a[x] < b[x] end end

--  _shuffle(list) --> list_   
-- `shuffle` does not change the input list. `
local function shuffle(t,    j)
  for i = #t, 2, -1 do j = R(i); t[i], t[j] = t[j], t[i] end
  return t end

-- _sort(list, callable=function(x) return x end) --> list_  
-- `sort` reorders the input list. 
local function sort(t,fn) 
  table.sort(t,fn); return t end

--- ### Mapping

-- _adds(Any, list) --> Any_
local function adds(it,t)
  for _,x in pairs(t or {}) do it:add(x) end; return it end

-- _kap(list, callable) --> list_   
-- `kap` calls `fn` with  key and value.
local function kap(t,fn,    u) --> list
  u={}; for k,v in pairs(t) do u[1+#u]=fn(k,v)  end; return u end

-- _map(list, callable) --> list_     
-- `kap` calls `fn` with  value.
local function map(t,fn,     u) --> list
  u={}; for _,v in pairs(t) do u[1+#u] = fn(v)  end; return u end

-- _sum(list,callable) --> num_
local function sum(t,fn,     n)
  n=0; for _,x in pairs(t) do n=n+(fn and fn(x) or x) end; return n end

-- _keysort(list, callable) --> list_    
-- Recommended over `sort` is `callable` is slow to compute.
local function keysort(t,fn,     DECORATE,UNDECORATE)
  DECORATE   = function(x) return {fn(x),x} end
  UNDECORATE = function(x) return x[2] end
  return map(sort(map(t,DECORATE),lt(1)), UNDECORATE) end

-- ### String to thing

-- `coerce(str) -> bool | int | float | str`
local function coerce(s,     F,TRIM) 
  TRIM= function(s) return s:match"^%s*(.-)%s*$" end
  F   = function(s) return s=="true" and true or s ~= "false" and s end
  return math.tointeger(s) or tonumber(s) or F(TRIM(s)) end

-- _csv(str) --> Iterator --> list_
local function csv(file,     src)
  if file and file ~="-" then src=io.input(file) end
  return function(     s,t)
    s = io.read()
    if   not s 
    then if src then io.close(src) end 
    else t={}; for s1 in s:gmatch"([^,]+)" do push(t,coerce(s1)) end; return t end end end

local function datas(      i)
  i=0
  return function()
    while i<#arg do
      i = i+1
      if arg[i]:find"csv" then
        return arg[i], DATA:new():read(arg[i]) end end end end

-- ### Thing to string

-- _o(atom | dict | list) --> str_
local function o(x,     F,G,GO) --> str
  F  = function() return #x>0 and map(x,o) or sort(kap(x,G)) end
  G  = function(k,v) if GO(k) then return fmt(":%s %s",k,o(x[k])) end end
  GO = function(k,v) return not o(k):find"^_" end
  return type(x)=="number" and fmt("%g",x) or  
         type(x)~="table"  and tostring(x) or 
         "{" .. table.concat(F()," ") .. "}" end 

-- _oo(atom | dict | list) --> nil_     
-- Prints the string generated by `o`.
local function oo(x) 
  print(o(x)) end

-- ### polymorphism

-- _new(list, list) --> list_
local function new(klass, obj)
  klass.__index    = klass
  klass.__tostring = klass.__tostring or o
  return setmetatable(obj, klass) end

-- ### Start-up

-- _cli(dict) --> dict__    
-- Update slot xxx in `dict` if there is a command line flag `-x`
local function cli(t)
  for k,v in pairs(t) do
    v = tostring(v)
    for n,x in ipairs(arg) do
      if x=="-"..(k:sub(1,1)) then
        v= v=="false" and "true" or v=="true" and "false" or arg[n+1] end end 
    t[k] = coerce(v) end 
  return t end

-- `tests(dict[str,callable], list[str])`
local function tests(eg,all,      fails,ok,msg) 
  fails = 0
  for _,x in pairs(all) do
    math.randomseed(the.rseed)
    ok,msg = xpcall(eg[x], debug.traceback, _)
    if   ok == false or msg == false 
    then print(fail.." FAIL on "..x); fails = fails + 1
    else print(pass.." PASS on "..x) end end 
  os.exit(fails) end

---------------- ----------------- ----------------- ----------------- -----------------
function SYM:new(txt,at) 
  return new(SYM, {at=at, txt=txt, n=0, has={}, most=0, mode=nil}) end

function SYM:add(x)
  if x~="?" then
    self.n = self.n + 1
    self.has[x] = 1 + (self.has[x] or 0) 
    if self.has[x] > self.most then self.most, self.mode = self.has[x], x end end end

function SYM:like(x,prior)
  return  ((self.has[x] or 0) + the.m*prior) / (self.n + the.m)  end

----------------- ----------------- ----------------- ----------------- -----------------  
function NUM:new(txt,at) 
  return new(NUM, {at=at, txt=txt, n=0, mu=0, m2=0, sd=0, lo=big, hi= -big,
                   goal = (txt or ""):find"-$" and 0 or 1}) end

function NUM:add(x)
  if x~="?" then
    self.n = self.n + 1
    local d = x - self.mu
    self.mu = self.mu + d / self.n
    self.m2 = self.m2 + d * (x - self.mu)
    self.sd = self.n < 2 and 0 or (self.m2/(self.n - 1))^.5
    if x > self.hi then self.hi = x end
    if x < self.lo then self.lo = x end end end

function NUM:like(x,_,      v,tmp)
  v = self.sd^2 + 1/big
  tmp = exp(-1*(x - self.mu)^2/(2*v)) / (2*pi*v) ^ 0.5
  return max(0,min(1, tmp + 1/big)) end
  
function NUM:norm(x)
  return x=="?" and x or (x - self.lo)/(self.hi - self.lo) end
       
function NUM:delta(other,      y,z,e)
  return abs(self.mu - other.mu) / ((1E-32 + self.sd^2/self.n + other.sd^2/other.n)^.5) end
             
----------------- ----------------- ----------------- ----------------- -----------------  
function COLS:new(names)
  local all,x,y = {},{},{}
  for at,txt in pairs(names) do 
    push(all, (txt:find"^[A-Z]" and NUM or SYM):new(txt,at))
    if not txt:find"X$" then
      push(txt:find"[!+-]$" and y or x, all[#all]) end end
  return new(COLS, {names=names, all=all, x=x, y=y}) end

----------------- ----------------- ----------------- ----------------- -----------------  
function DATA:new()
  return new(DATA, {cols=nil, rows={}})  end

function DATA:read(src)
  for row in csv(src) do self:add(row) end; return self end

function DATA:adds(t)
  for row in pairs(t or {}) do self:add(row) end; return self end

function DATA:add(row) 
  if self.cols then
    push(self.rows, row)
    for _,col in pairs(self.cols.all) do col:add(row[col.at]) end 
	else
	  self.cols= COLS:new(row) end end 

function DATA:clone(rows) 
  return adds(adds(DATA:new(), {self.cols.names}),rows) end
    
function DATA:loglike(row, nall, nh)
  local prior,out,tmp
  prior = (#self.rows + the.k) / (nall + the.k*nh)
  out,tmp = 0,log(prior)
  for _,col in pairs(self.cols.y) do 
    tmp = col:like(row[col.at], prior)
    if tmp > 0 then
       out = out + log(tmp) end end
  return out end

function DATA:ydist(row, D)
  D = function(y) return (abs(y:norm(row[y.at]) - y.goal))^the.p end
  return (sum(self.cols.y,D) /#self.cols.y)^(1/the.p) end

-- _DATA:active() --> list,list_
function DATA:acquire()
  local Y,B,R,BR,test,train,todo,done,best,rest
  Y          = function(r) return self:ydist(r) end
  B          = function(r) return best:loglike(r, #done, 2) end
  R          = function(r) return rest:loglike(r, #done, 2) end
  BR         = function(r) return B(r) - R(r) end
  test,train = split(shuffle(self.rows), (the.Holdout * #self.rows))
  done,todo  = split(train, the.start)
  while true do
    done = keysort(done,Y)
    if #done > the.Stop or #todo < 5 then break end
    best,rest = split(done, sqrt(#done))
    best,rest = self:clone(best), self:clone(rest)
    todo      = keysort(todo,BR)
    for _=1,2 do push(done, pop(todo));   push(done, pop(todo,1)) end
  end
  return done[1], keysort(test,BR)[#test] end

----------------- ----------------- ----------------- ----------------- -----------------  
-- ## Stats

-- _cliffs(list[num], list[num]) --> bool_
local function cliffs(xs,ys)
  local lt,gt,n = 0,0,0
  for _,x in pairs(xs) do
     for _,y in pairs(ys) do
       n = n + 1
       if y > x then gt = gt + 1 end
       if y < x then lt = lt + 1 end end end
  return abs(gt - lt)/n <= the.Cliffs end -- 0.195 

-- _boostrap(list[num], list[num]) --> bool_ 
local function bootstrap(y0,z0)
  --non-parametric significance test From Introduction to Bootstrap,
  -- Efron and Tibshirani, 1993, chapter 20. https://doi.org/10.1201/9780429246593
  local x,y,z,delta0,yhat,zhat,n,this,that,b
  z,y,x  = adds(NUM:new(),z0), adds(NUM:new(),y0), adds(adds(NUM:new(),y0),z0)
  delta0 = y:delta(z)
  yhat   = map(y0, function(y1) return y1 - y.mu + x.mu end)
  zhat   = map(z0, function(z1) return z1 - z.mu + x.mu end)
  n,b    = 0, the.Bootstraps
  for i=1, b do
    this, that = adds(NUM:new(),many(yhat)), adds(NUM:new(),many(zhat))
    n = n + (this:delta(that) > delta0 and 1 or 0) end
  return n / b >= the.conf end

---------------- ----------------- ----------------- ----------------- -----------------  
local EG={}

function EG.all() 
  return tests(EG, {"any","sum","split","sort","num","sym","csv"}) end

function EG.any(  a)
  a = {10,20,30,40,50,60}
	oo(map({any(a), many(a,3),shuffle(a),keysort(a,function(x) return -x end)},
     o)) end 

function EG.sum()
  assert(210 == sum{10,20,30,40,50,60}) end

function EG.split(    a,b,c)
  a={10,20,30,40,50,60}
	b,c = split(a,3)
	print(o(b), o(c)) end

function EG.sort(    t)
  t={1,2,3,4,5,6,7}
  t=sort(t, function(x,y) return  x > y end)
  oo{10,4,5}
  oo(t) end

function EG.num(    n) 
  n = NUM:new()
  for _ = 1,1000 do n:add( normal(10,2) ) end
  assert(10-n.mu < 0.1 and 2-n.sd < 0.03) end

function EG.sym(    s) 
  s = adds(SYM:new(), {"a","a","a","a","b","b","c"})
	print(s.mode, o(s.has)) end

function EG.csv(   d, n)
  print(the.train)
  n=0
  for row in csv(the.train) do n=n+1 ; if n==1 or n % 90==0 then oo(row) end end end

function EG.stats(   r,t,u,d)
  d,r= 1,100
  while d< 1.2 do
    t={}; for i=1,r do t[1+#t] = normal(10,2)^2 end 
    u={}; for i,x in pairs(t) do u[i] = x*d end
    d=d*1.01
    print(fmt("%.3f    %s %s",d, cliffs(t,u) and "y" or ".", bootstrap(t,u) and "y" or "."))
  end end

function EG.data(   d)
  d = DATA:new():read(the.train) 
  assert(3184 == (sum(d.rows,function(r) return #r end))) end

function EG.ydist(   d,n)
  n,d = 0,DATA:new():read(the.train) 
  for n,row in pairs(keysort(d.rows, function(r) return d:ydist(r) end)) do
     if n==1 or n % 30==0 then print(n,o(row), d:ydist(row)) end end end

function EG.like(   d,n)
  n,d = 0,DATA:new():read(the.train) 
  for _,row in pairs(d.rows) do 
    n=n+1 
    if n==1 or n % 15==0 then 
      print(fmt("%.3f %s",d:loglike(row,#d.rows,2), o(row))) end end  end

function EG.acquire()
  d = DATA:new():read(the.train) 
  train,test = d:acquire() 
  print(d:ydist(train), d:ydist(test)) end

local SOME={}
function SOME:new(txt, items)
  return adds(new(SOME, {txt=txt, all={}, num=NUM:new()}), items) end

function SOME:add(x)
  self.num:add(x)
  push(self.all, x) end 

function SOME:diff(other,eps=0)
  if cliffs(y,z) and bootstrap(y,z) then return 0 end
  pooled = sqrt(((self.n-1)*self.sd^2 + (other.n-1)*other.sd^2) / (self.n + other.n-2))
  d = self.all.num.mu - other.all.num.mu
  return abs(d) < pooled*eps and 0 or d end

function EG.acquires(     d,y,trains,tests,train,test,r,asIs,num0,num1,num2,eps,diff)
  r = 20
  for file,d in datas(arg) do
    Y= function(r) return d:ydist(r) end
    asIs = SOME:new("asIs", map(d.rows, Y))
    for _,n in pairs{15,30,50,80,120} do
      the.Stop=n
      trains,tests = SOME:new("train"), SOME:new("test")
      for i=1, r do
        train,test = d:acquire() 
        trains:add(Y(train)) 
        tests:add(Y(test)) end
      diff=train:diff(test)
      oo{file=file:gsub(".*/",""), n=the.Stop, eps=eps,mu0=asIs.mu, 
         mu1=num1.mu, mu2=num2.mu, diff=diff}
      
----------------- ----------------- ----------------- ----------------- -----------------  
help:gsub("%s+-%S%s(%S+)[^=]+=%s+(%S+)%s*\n", function(k,v) the[k]= coerce(v) end)

if   o(arg):find"kah.lua" 
then the = cli(the)
     if the.help then os.exit(print(help)) end
     math.randomseed(the.rseed or 1)
     map(arg, function(s) if EG[s:sub(3)] then EG[s:sub(3)]() end end) end 

return {SYM=SYM, NUM=NUM, COLS=COLS, DATA=DATA, the=the, help=help}
