#!/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


# 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,y,z = numpy.loadtxt('df_in.csv',delimiter=',', unpack=True)
t = t[:,None] # make column vector
y = y[:,None]
z = z[:,None]

n = y.shape[0]
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

x=None
t_vis = numpy.loadtxt('df_vis.csv',delimiter=',')

n_vis=t_vis.shape[0]
x_vis=None

        
print("Data shapes:\n\
Features:{x},\n\
Instruments:{z},\n\
Treament:{t},\n\
Response:{y}".format(**{'x':x.shape if x is not None else None, 'z':z.shape,
                        't':t.shape, 'y':y.shape}))

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

hidden = [128, 64, 32]

act = "relu"

n_components = 10

est_treat = architectures.feed_forward_net(instruments, 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], outputs=est_treat)
treatment_model.compile('adam',
                        loss="mixture_of_gaussians",
                        n_components=n_components)

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

# Build and fit response model

treatment = Input(shape=(t.shape[1],), name="treatment")

est_response = architectures.feed_forward_net(treatment, 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=[treatment],
                            outputs=est_response)
    response_model.compile('adam', loss='mse', unbiased_gradient=unbiased, batch_size=batch_size)
    response_model.fit([z], y, epochs=epochs, verbose=1,
                    batch_size=batch_size, samples_per_batch=2)


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