#!/anaconda2/envs/deepIV/bin/python
# -*- coding: utf-8 -*-

from __future__ import print_function

activate_this_file = "/anaconda2/envs/deepIV/bin/activate_this.py"

execfile(activate_this_file, dict(__file__=activate_this_file))

#import warnings

from deepiv.models import Treatment, Response
import deepiv.architectures as architectures
import deepiv.densities as densities

#import tensorflow as tf

from keras.layers import Input, Dense
#from keras.models import Model
from keras.layers.merge import Concatenate

import numpy
#import data_generator
import csv


#def datafunction(n, s, images=images, test=False):
#    return data_generator.demand(n=n, seed=s, ypcor=0.5, use_images=images, test=test)
#
#x, z, t, y, g_true = datafunction(n, 1)
#
#print("Data shapes:\n\
#Features:{x},\n\
#Instruments:{z},\n\
#Treament:{t},\n\
#Response:{y}".format(**{'x':x.shape, 'z':z.shape,
#                        't':t.shape, 'y':y.shape}))

# SSG (x,y,z) maps to HLLT code (t,y,z)
# SSG (p,(t,s),y,z) maps to HLLT code (t,x,y,z)
t=[]
x1=[]
x2=[]
y=[]
z=[]


with open('df_in.csv') as csvfile:
    readCSV = csv.reader(csvfile, delimiter=',')
    for row in readCSV:
        t.append(row[0])
        x1.append(row[1])
        x2.append(row[2])
        y.append(row[3])
        z.append(row[4])
csvfile.close()


n = len(y)
dropout_rate = min(1000./(1000. + n), 0.5)
epochs = int(1500000./float(n)) # heuristic to select number of epochs
epochs = 300
batch_size = 100
images = False
        
t=numpy.asarray(t)
t.shape=(n,1)

x1=numpy.asarray(x1)
x1.shape=(n,1)

x2=numpy.asarray(x2)
x2.shape=(n,1)

x=numpy.concatenate((x1,x2),axis=1)

y=numpy.asarray(y)
y.shape=(n,1)

z=numpy.asarray(z)
z.shape=(n,1)

t_vis=[]
x1_vis=[]
x2_vis=[]
with open('df_vis.csv') as csvfile:
    readCSV = csv.reader(csvfile, delimiter=',')
    for row in readCSV:
        t_vis.append(row[0])
        x1_vis.append(row[1])
        x2_vis.append(row[2])
csvfile.close()

n_vis=len(t_vis)

t_vis=numpy.asarray(t_vis)
t_vis.shape=(n_vis,1)

x1_vis=numpy.asarray(x1_vis)
x1_vis.shape=(n_vis,1)

x2_vis=numpy.asarray(x2_vis)
x2_vis.shape=(n_vis,1)
x_vis=numpy.concatenate((x1_vis,x2_vis),axis=1)
        
print("Data shapes:\n\
Features:{x},\n\
Instruments:{z},\n\
Treament:{t},\n\
Response:{y}".format(**{'x':x.shape, 'z':z.shape,
                        't':t.shape, 'y':y.shape}))




# Build and fit treatment model
instruments = Input(shape=(z.shape[1],), name="instruments")
features = Input(shape=(x.shape[1],), name="features")
treatment_input = Concatenate(axis=1)([instruments, features])

hidden = [128, 64, 32]

act = "relu"

n_components = 10

est_treat = architectures.feed_forward_net(treatment_input, lambda x: densities.mixture_of_gaussian_output(x, n_components),
                                           hidden_layers=hidden,
                                           dropout_rate=dropout_rate, l2=0.0001,
                                           activations=act)

treatment_model = Treatment(inputs=[instruments, features], outputs=est_treat)
treatment_model.compile('adam',
                        loss="mixture_of_gaussians",
                        n_components=n_components)

treatment_model.fit([z, x], t, epochs=epochs, batch_size=batch_size)

# Build and fit response model

treatment = Input(shape=(t.shape[1],), name="treatment")
response_input = Concatenate(axis=1)([features, treatment])

est_response = architectures.feed_forward_net(response_input, Dense(1),
                                              activations=act,
                                              hidden_layers=hidden,
                                              l2=0.001,
                                              dropout_rate=dropout_rate)



for unbiased in [True, False]:
    response_model = Response(treatment=treatment_model,
                            inputs=[features, treatment],
                            outputs=est_response)
    response_model.compile('adam', loss='mse', unbiased_gradient=unbiased, batch_size=batch_size)
    response_model.fit([z, x], y, epochs=epochs, verbose=1,
                    batch_size=batch_size, samples_per_batch=2)


    y_pred=response_model.predict([x_vis, t_vis])
    numpy.savetxt("df_out_%s.csv" % ('unbiased' if unbiased else 'biased'), y_pred, delimiter=",",fmt='%f')

#response_model = Response(treatment=treatment_model,
#                          inputs=[features, treatment],
#                          outputs=est_response)
#response_model.compile('adam', loss='mse')
#response_model.fit([z, x], y, epochs=epochs, verbose=1,
#                   batch_size=batch_size, samples_per_batch=2)
#
#
#y_pred=response_model.predict([x_vis,t_vis])
#numpy.savetxt("df_out.csv", y_pred, delimiter=",")

#oos_perf = data_generator.monte_carlo_error(lambda x,z,t: response_model.predict([x,t]), datafunction, has_latent=images, debug=False)
#print("Out of sample performance evaluated against the true function: %f" % oos_perf)
