local xlua = require 'xlua'
local pl_data = require 'pl.data'
local pl_string = require 'pl.stringx'
local rk = require 'randomkit'


local function matrix2file(matrix, fname)

  local out = io.open(fname, "w")

  for i=1,matrix:size(1) do

      for j=1,matrix:size(2) do

          out:write(matrix[i][j])

          if j == matrix:size(2) then

              out:write("\n")

          else

              out:write(" ")

          end

      end

  end

  out:close()

end



local function normalize(x)

  local y = x:clone()

  local m = x:mean(1):mul(-1)

  local s = x:std(1):add(1e-3)

  y:add(m:expandAs(y)):cdiv(s:expandAs(y))

  return y

end



local function gaussian(m)

  local m = m or 2600

  local a = torch.randn(m)[1]

  local b = torch.randn(m)[1]

  return torch.randn(m,1):mul(a):add(b)

end



local function uniform(m)

  local m = m or 2600

  local a = torch.randn(m)[1]

  local b = torch.randn(m)[1]

  return normalize(torch.rand(m,1)):mul(a):add(b)

end



local function linear(x)

  local a = torch.randn(1)[1]

  local b = torch.randn(1)[1]

  foo = function(x)

    local f = x:clone()

    return f:mul(a):add(b)

  end

  return foo

end



local function parabola(x)

  local a = torch.randn(1)[1]

  local b = torch.randn(1)[1]

  foo = function(x)

    local f = torch.pow(x,2) 

    return f:mul(a):add(b)

  end

  return foo

end



local function post_additive(x,f,n)

  return f(x):add(n)

end



local function post_multiplicative(x,f,n)

  return f(x):cmul(n)

end



local function pre_additive(x,f,n)

  return f(torch.add(x,n))

end



local function pre_multiplicative(x,f,n)

  return f(torch.cmul(x,n))

end



local marginals = {

  gaussian = gaussian,

  uniform = uniform

}

local tab_ma={"gaussian","uniform"}



    

local mechanisms = {

  linear = linear,

  parabola = parabola

}

local tab_me = {

    "linear",

    "parabola"

}

local combinations = {

  post_additive = post_additive,

  post_multiplicative = post_multiplicative,

  pre_additive = pre_additive,

  pre_multiplicative = pre_multiplicative

}

local tab_co={

    "post_additive",

    "post_multiplicative",

    "pre_additive",

    "pre_multiplicative"

}



local function draw_rand_parameters()

    cause_j = tab_ma[torch.Tensor(1):random(1,#tab_ma)[1]]

    noise_j = tab_ma[torch.Tensor(1):random(1,#tab_ma)[1]]

    mechanism_j = tab_me[torch.Tensor(1):random(1,#tab_me)[1]]

    combination_j = tab_co[torch.Tensor(1):random(1,#tab_co)[1]]

    

    return cause_j,noise_j,mechanism_j,combination_j

end



local function pair(cause, mechanism, noise, combination)

  local x = marginals[cause]()

  local f = mechanisms[mechanism]()

  local n = marginals[noise]()

  local y = combinations[combination](x,f,n)

  return torch.cat(normalize(x), normalize(y), 2)

   

end





local function identifiable(cause, mechanism, noise, combination)

  if((cause == "gaussian") and

     (mechanism == "linear") and

     (noise == "gaussian")) then

     if((combination == "post_additive") or

        (combination == "pre_additive")) then

        return false

     end

  end

  

  return true

end







local function r_pair(typep)

    

  local cause_l,noise_l,mechanism_l,combination_l=draw_rand_parameters()  

  if(identifiable(cause_l, mechanism_l, noise_l, combination_l)) then

    --print(cause_l)

    local x = marginals[cause_l]()

    local f = mechanisms[mechanism_l]()

    local n = marginals[noise_l]()

    local y = combinations[combination_l](x,f,n)

    

    if typep==1 then

        return torch.cat(normalize(x), normalize(y), 2)

    elseif typep==2 then

        return torch.cat(normalize(y), normalize(x), 2)

    elseif typep==3 then

        local cause_l,noise_l,mechanism_l,combination_l=draw_rand_parameters()  

        if(identifiable(cause_l, mechanism_l, noise_l, combination_l)) then

            local f = mechanisms[mechanism_l]()

            local n = marginals[noise_l]()

            local z = combinations[combination_l](x,f,n)

            return torch.cat(normalize(z), normalize(y), 2)

        else 

            return r_pair(typep)

        end

    else

       local cause_l,noise_l,mechanism_l,combination_l=draw_rand_parameters() 

       local x = marginals[cause_l]()

       local f = mechanisms[mechanism_l]()

       local n = marginals[noise_l]()

       local y2 = combinations[combination_l](x,f,n)

       return torch.cat(normalize(y2), normalize(y), 2)

    end

  else 

    return r_pair(typep)

  end    

end





j=1

nb_pairs= 4000

local meta = io.open('pairmeta.txt', "w")

while j < nb_pairs do

    --X->Y

      pair_i = r_pair(1)

      matrix2file(pair_i, "pairF" .. j .. ".txt")

      meta:write("pairF" .. j .. ' ' .. '1\n')

      j=j+1

      

      xlua.progress(j,nb_pairs)

      

    --X<-Y

    

      pair_i = r_pair(2)

      matrix2file(pair_i, "pairF" .. j .. ".txt")

      meta:write("pairF" .. j .. ' ' .. '2\n')

      j=j+1

    

      xlua.progress(j,nb_pairs)

    

    --X||Y

    

      pair_i = r_pair(3)  

      matrix2file(pair_i, "pairF" .. j .. ".txt")

      meta:write("pairF" .. j .. ' ' .. '3\n')

      j=j+1

      xlua.progress(j,nb_pairs)

      

    --X_|_Y

    

      pair_i = r_pair(4)

      matrix2file(pair_i, "pairF" .. j .. ".txt")

      meta:write("pairF" .. j .. ' ' .. '4\n')

      j=j+1

      xlua.progress(j,nb_pairs)

      

end

meta:close()



