"""Defines the neural network, losss function and metrics"""

import os

import numpy as np

import pandas as pd

import torch

import torch.nn as nn

import torch.nn.functional as F

from torch.distributions.multivariate_normal import MultivariateNormal

from torchvision import models

from scipy.stats import spearmanr

import matplotlib.pyplot as plt

import seaborn as sns

class Flatten(nn.Module):

    def forward(self, input):

        return input.view(input.size(0), -1)

class Identity(nn.Module):

    def __init__(self, *args, **kwargs):

        super(Identity, self).__init__()

    def forward(self, x):

        return x

class ConcatRegressor(nn.Module):

    """

    Module for concatenating feature info in final layer

    Always includes conditioning on t, optionally on x

    """

    def __init__(self, in_features=144, concat_dim=1):

        super(ConcatRegressor, self).__init__()

        self.fc = nn.Linear(in_features, 1)

        self.t  = nn.Linear(concat_dim, 1, bias=False)

        nn.init.constant_(self.t.weight, 0)

    def forward(self, x, t):

        return self.fc(x) + self.t(t)

class SimpleEncoder(nn.Module):

    def __init__(self, params, setting):

        super(SimpleEncoder,self).__init__()

        self.params = params

        self.setting =  setting

        self.fwd = nn.Sequential(

            nn.Conv2d(1, 16, 3, stride=1, padding=1),

            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            nn.Conv2d(16, 16, 3, stride=1, padding=1),

            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            nn.Conv2d(16, 16, 3, stride=1, padding=1),

            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            nn.Conv2d(16, 16, 3, stride=1, padding=1),

            nn.ReLU(inplace=True),

            nn.MaxPool2d(2),

            nn.Conv2d(16, 16, 3, stride=1, padding=1),

            nn.ReLU(inplace=True),

            nn.AvgPool2d((1,1)),

            Flatten()

        )

    def forward(self, x, t=None):

        return self.fwd(x)

encoders = {'simple': SimpleEncoder}

class CausalNet(nn.Module):

    def __init__(self, params, setting):

        super(CausalNet, self).__init__()

        self.params  = params

        self.setting = setting

        # storage for betas from OLS

        # keep in model to port from train to valid

        self.betas_bias   = torch.zeros((params.regressor_z_dim+2,1), requires_grad=False) 

        self.betas_causal = torch.zeros((params.regressor_z_dim+1,1), requires_grad=False)

        print("instantiating net")

        self.encoder = encoders[setting.encoder](params, setting)

        if setting.encoder == 'simple':

            fc_in_features = 144

        else:

            raise NotImplementedError(f'different encoder than simple currently not implemented: {setting.encoder})')

        # pick the right type of regressor, possibly allowing for interactions

        if params.conditioning_place == "regressor":

            Regressor = ConcatRegressor

        else:

            raise NotImplementedError('only conditioning in final layer is implemented now')

        # same size in and out fcs

        self.fcs = nn.ModuleList(params.num_fc*[

            nn.Linear(fc_in_features, fc_in_features), 

            nn.ReLU(inplace=True),

            nn.Dropout(params.dropout_rate)

            ])

        # fc layer to final regression layer

        # NOTE keep track if a ReLU is needed here (probably not)

        self.fcr = nn.Linear(fc_in_features, params.regressor_z_dim + params.regressor_x_dim)

        # final regressor to y; this takes in entire last layer and treatment

        self.regressor = Regressor(params.regressor_z_dim+params.regressor_x_dim, concat_dim=1)

        # initialize weights

        for layer_group in [self.encoder, self.fcs, self.fcr, self.regressor]:

            for module in layer_group.modules():

                if hasattr(module, 'weight'):

                    torch.nn.init.xavier_uniform_(module.weight)

    def forward(self, x, t=None, epoch=None):

        # prepare dictionary for keeping track of output tensors

        outs = {}

        # convolutional stage to get 'features'

        h = self.encoder(x)

        # pass through a sequence of same-size in-out fc-layers for 'non-linear interactions'

        for i, module in enumerate(self.fcs):

            h = module(h)

        # squeeze to lower size for final regression layer

        finalactivations = self.fcr(h)

        # store tensors ('bottlenecks' from which correlations / MIs are calculated)

        outs['bnx'] = finalactivations[:,:self.params.regressor_x_dim] # activations that represent x

        outs['bnz'] = finalactivations[:,self.params.regressor_x_dim:] # activations that represent z (=everything else)

        # predict y from final activations and treatment

        outs['y'] = self.regressor(finalactivations, t)

        return outs

def freeze_conv_layers(model, keep_layers = ["bnx", "bny", "bnbnx", "bnbny", "fcx", "fcy", "t"], last_frozen_layer=None):

    for name, param in model.named_parameters():

        if name.split(".")[0] not in keep_layers:

            param.requires_grad = False

        else:

            print("keeping grad on for parameter {}".format(name))

def speedup_t(model, params):

    lr_t = params.lr_t_factor * params.learning_rate

    optimizer = torch.optim.Adam(model.regressor.t.parameters(), lr = lr_t)

    if params.speedup_intercept:

        optimizer.add_param_group({'params': model.regressor.fc.bias, 'lr': lr_t})

    for name, param in model.named_parameters():

        # print(f"parameter name: {name}")

        if name.split(".")[1] == "t":

            print("Using custom lr for param: {}".format(name))

        elif name.endswith("fc.bias") and params.speedup_intercept:

            print("Using cudtom lr for param: {}".format(name))

        else:

            optimizer.add_param_group({'params': param, 'lr': params.learning_rate, 'weight_decay': params.wd})

    return optimizer

def softfreeze_conv_layers(model, params, fast_layers = ["bnx", "bny", "bnbnx", "bnbny", "fcx", "fcy"], last_frozen_layer=None):

    optimizer = torch.optim.Adam(model.t.parameters(), lr=params.learning_rate)

    for name, param in model.named_parameters():

        if name in fast_layers:

            optimizer.add_param_group({'params': param})

        elif name.split(".")[0] == "t":

            pass

        else:

            optimizer.add_param_group({'params': param, 'lr': params.learning_rate / 10})

    return optimizer

def get_loss_fn(setting, **kwargs):

    if setting.num_classes == 2:

        print("Loss: cross-entropy")

        def loss_fn(outputs, labels, **kwargs):

            criterion = nn.CrossEntropyLoss(**kwargs)

            target = labels.type(torch.cuda.LongTensor)

            # print(target.size())

            # print(outputs.size())

            return criterion(outputs, target)

    else: 

        print("Loss: MSE")

        def loss_fn(outputs, labels, **kwargs):

            criterion = nn.MSELoss()

            # return torch.sqrt(criterion(outputs.squeeze(), labels.squeeze()))

            return criterion(outputs.squeeze(), labels.squeeze())

    return loss_fn

def bottleneck_loss(bottleneck_features):

    z_mean    = bottleneck_features, outputs, labels

    z_stddev  = bottleneck_features, outputs, labels

    mean_sq   = z_mean * z_mean

    stddev_sq = z_stddev * z_stddev, outputs, labels

    return 0.5 * torch.mean(mean_sq + stddev_sq - torch.log(stddev_sq + 1.0e-6) - 1)

def get_bn_loss_fn(params):

    if params.bn_loss_type == "variational-gaussian":

        def loss_fn(outputs):

            # take mean and sd over batch dimension

            z_mean    = outputs.mean(0)

            z_stddev  = outputs.std(0)

            mean_sq   = z_mean * z_mean

            stddev_sq = z_stddev * z_stddev

            return 0.5 * torch.mean(mean_sq + stddev_sq - torch.log(stddev_sq + 1.0e-6) - 1)

    else:

        raise NotImplementedError

    return loss_fn

def rmse(setting, model, outputs, labels, data=None):

    return np.sqrt(np.mean(np.power((outputs - labels), 2)))

def bias(setting, model, outputs, labels, data=None):

    weights = model.t.weight.detach().cpu().numpy()

    return np.squeeze(weights)[-1] - 1

def b_t(setting, model, outputs, labels, data=None):

    weight = model.regressor.t.weight.detach().cpu().numpy()

    return weight

def intercept(setting, model, outputs, labels, data=None):

    # oracle = pd.read_csv(os.path.join(setting.data_dir, "oracle.csv"))

    bias = model.cnn.fc2.bias.detach().cpu().numpy()

    return bias

    # for now: use ATE = 1

def ate(setting, model, outputs, labels, data):

    # data should always have treatment in first columns

    if data.ndim == 1:

        t = data

    else:

        t = data[:,0].squeeze()

    treated   = outputs[np.where(t)]

    untreated = outputs[np.where(t == 0)]

    return treated.mean() - untreated.mean()

def total_loss(setting, model, outputs, labels, data=None):

    # total_loss_fn = get_loss_fn(setting, reduction="sum")

    total_loss_fn = nn.MSELoss(reduction="sum")

    outputs = torch.tensor(outputs, requires_grad=False).squeeze()

    labels  = torch.tensor(labels, requires_grad=False).squeeze()

    return total_loss_fn(outputs, labels)

def accuracy(setting, model, outputs, labels, data=None):

    """

    Compute the accuracy, given the outputs and labels for all images.

    Args:

        outputs: (np.ndarray) dimension batch_size x 6 - log softmax output of the model

        labels: (np.ndarray) dimension batch_size, where each element is a value in [0, 1, 2, 3, 4, 5]

    Returns: (float) accuracy in [0,1]

    """

    outputs = np.argmax(outputs, axis=1)

    return np.sum(outputs==labels)/float(labels.size)

def ppv(setting, model, outputs, labels, data=None):

    if setting.num_classes == 2:

        pos_preds = np.argmax(outputs, axis=1)==1

        if pos_preds.sum() > 0:

            return accuracy(setting, model, outputs[pos_preds,:], labels[pos_preds])

        else:

            return np.nan

    else:

        return 0.

def npv(setting, model, outputs, labels, data=None):

    if setting.num_classes == 2:

        neg_preds = np.argmax(outputs, axis=1)==0

        if neg_preds.sum() > 0:

            return accuracy(setting, model, outputs[neg_preds,:], labels[neg_preds])

        else:

            return np.nan

    else:

        return 0.

def cholesky_least_squares(X, Y, intercept=True):

    """

    Perform least squares regression with cholesky decomposition 

    intercept: add intercept to X

    adapted from https://gist.github.com/gngdb/611d8f180ef0f0baddaa539e29a4200e

    which was adapted from http://drsfenner.org/blog/2015/12/three-paths-to-least-squares-linear-regression/

    """

    if X.ndimension() == 1:

        X.unsqueeze_(1)    

    if intercept:

        X = torch.cat([torch.ones_like(X[:,0].unsqueeze(1)),X], dim=1)

    XtX, XtY = X.permute(1,0).mm(X), X.permute(1,0).mm(Y)

    betas, _ = torch.gesv(XtY, XtX)

    return betas.squeeze()

def mse_loss(output, target):

    criterion = nn.MSELoss()

    return criterion(output, target)

def spearmanrho(outputs, labels):

    '''

    calculate spearman (non-parametric) rank statistic

    '''

    try:

        return spearmanr(outputs.squeeze(), labels.squeeze())[0]

    except ValueError:

        print('value error in spearmanr, returning 0')

        return np.array(0)

# maintain all metrics required in this dictionary- these are used in the training and evaluation loops

all_metrics = {

    'total_loss': total_loss,

    'bottleneck_loss': bottleneck_loss,

    'accuracy': accuracy,

    'rmse': rmse,

    'bias': bias,

    'ate': ate,

    'intercept': intercept,

    'b_t': b_t,

    'ppv': ppv,

    'npv': npv,

    'spearmanrho': spearmanrho

    # 'ite_mean': ite_mean

    # could add more metrics such as accuracy for each token type

}

# from here: https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2

def cov(m, rowvar=False):

    '''Estimate a covariance matrix given data.

    Covariance indicates the level to which two variables vary together.

    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,

    then the covariance matrix element `C_{ij}` is the covariance of

    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.

    Args:

        m: A 1-D or 2-D array containing multiple variables and observations.

            Each row of `m` represents a variable, and each column a single

            observation of all those variables.

        rowvar: If `rowvar` is True, then each row represents a

            variable, with observations in the columns. Otherwise, the

            relationship is transposed: each column represents a variable,

            while the rows contain observations.

    Returns:

        The covariance matrix of the variables.

    '''

    if m.dim() > 2:

        raise ValueError('m has more than 2 dimensions')

    if m.dim() < 2:

        m = m.view(1, -1)

    if not rowvar and m.size(0) != 1:

        m = m.t()

    # m = m.type(torch.double)  # uncomment this line if desired

    fact = 1.0 / (m.size(1) - 1)

    m -= torch.mean(m, dim=1, keepdim=True)

    mt = m.t()  # if complex: mt = m.t().conj()

    return fact * m.matmul(mt).squeeze()

def get_of_diag(x):

    '''

    Set the diagonal elements of a matrix to zero, and flatten the rest

    '''

    assert type(x) is np.ndarray

    x = x[~np.eye(x.shape[0],dtype=bool)]

    return x.reshape(-1,1)

def make_scatter_plot(x,y,c=None,

                      xlabel: str=None,ylabel: str=None,title: str=None):

    '''

    make scatter plots for tensorboard

    '''

    if c is not None:

        g = sns.jointplot(x.reshape(-1,1),y.reshape(-1,1), kind='reg')

        # g = sns.jointplot(x.reshape(-1,1),y.reshape(-1,1), joint_kws=dict(scatter_kws=dict(c=c.reshape(-1,1))), kind='reg')

    else:

        g = sns.jointplot(x.reshape(-1,1),y.reshape(-1,1), kind='reg')

    g.set_axis_labels(xlabel, ylabel)

    g.ax_joint.set_title(xlabel+ " vs " + ylabel)

    return g.fig