from typing import Any, Callable, List

import numpy as np

import torch

from torch import nn

import os

import tqdm

import catenets.logger as log

from catenets.models.constants import (

    DEFAULT_BATCH_SIZE,

    DEFAULT_DIM_P_OUT,

    DEFAULT_DIM_P_R,

    DEFAULT_DIM_S_OUT,

    DEFAULT_DIM_S_R,

    DEFAULT_LAYERS_OUT,

    DEFAULT_LAYERS_R,

    DEFAULT_N_ITER,

    DEFAULT_N_ITER_MIN,

    DEFAULT_N_ITER_PRINT,

    DEFAULT_PATIENCE,

    DEFAULT_PENALTY_L2,

    DEFAULT_PENALTY_ORTHOGONAL,

    DEFAULT_SEED,

    DEFAULT_NJOBS,

    DEFAULT_STEP_SIZE,

    DEFAULT_VAL_SPLIT,

    LARGE_VAL,

)

from catenets.models.torch.base import DEVICE, BaseCATEEstimator

from catenets.models.torch.utils.model_utils import make_val_split

import pandas as pd

# Hydra

from omegaconf import DictConfig

import json

import datetime

from .src.main_model_table import TabCSDI

from .src.utils_table import train

from .dataset_acic import get_dataloader

from .PropensityNet import load_data

torch.manual_seed(0)

class AverageMeter(object):

    """Computes and stores the average and current value"""

    def __init__(self):

        self.reset()

    def reset(self):

        self.val = 0

        self.avg = 0

        self.sum = 0

        self.count = 0

    def update(self, val, n=1):

        self.val = val

        self.sum += val * n

        self.count += n

        self.avg = self.sum / self.count

class DiffPOLearner(BaseCATEEstimator):

    """

    A flexible treatment effect estimator based on the EconML framework.

    """

    def __init__(

        self,

        cfg: DictConfig,

        num_features: int,

        binary_y: bool,

    ) -> None:

        self.config = cfg.DiffPOLearner

        self.diffpo_path = cfg.diffpo_path

        self.config.diffusion.cond_dim = num_features+1 # make sure inner dimension matches the dataset

        self.est = None

        self.propnet = None

        self.device = DEVICE

        self.cate_cis = None # confidence intervals, dim: 2, n, num_T-1, dim_Y

        self.pred_outcomes = None

        # create folder if diffpo_path + 'data' does not exist

        if not os.path.exists(self.diffpo_path):

            os.makedirs(self.diffpo_path)

        # Store data for their pipeline

        self.data_dir = self.diffpo_path+'/data/'

        if not os.path.exists(self.data_dir):

            os.makedirs(self.data_dir)

        return None

    def reshape_data(self, X: np.ndarray, w: np.ndarray, outcomes: np.ndarray) -> None:

        data = np.concatenate([w.reshape(-1,1),outcomes[:,0],outcomes[:,1],outcomes[:,0],outcomes[:,1],X], axis=1)

        data_df = pd.DataFrame(data)

        # Create masking array of same shape as pp_data and initialize with 1s

        mask = np.ones(data_df.shape)

        mask[:,1] = w

        mask[:,2] = 1-w

        mask[:,3] = 0

        mask[:,4] = 0

        mask_df = pd.DataFrame(mask)

        return data_df, mask_df

    def train(self, X: np.ndarray, y: np.ndarray, w: np.ndarray, outcomes:np.ndarray) -> None:

        """

        Prepare data and train DiffPO Learner

        """

        log.info("Training data shapes: X: {}, Y: {}, T: {}".format(X.shape, y.shape, w.shape))

        if not os.path.exists(self.data_dir):

            os.makedirs(self.data_dir)

        data, mask = self.reshape_data(X, w, outcomes)

        # create destination folders if not exist

        if not os.path.exists(self.data_dir+"acic2018_norm_data/"):

            os.makedirs(self.data_dir+"acic2018_norm_data/")

        if not os.path.exists(self.data_dir+"acic2018_mask/"):

            os.makedirs(self.data_dir+"acic2018_mask/")

        # save intermediate data

        data.to_csv(self.data_dir+"acic2018_norm_data/data_pp.csv", index=False)

        mask.to_csv(self.data_dir+"acic2018_mask/data_pp.csv", index=False)

        # Remove old files

        if os.path.exists(self.data_dir+"missing_ratio-0.2_seed-1_current_id-data_max-min_norm.pk"):

            os.remove(self.data_dir+"missing_ratio-0.2_seed-1_current_id-data_max-min_norm.pk")

        if os.path.exists(self.data_dir+"missing_ratio-0.2_seed-1.pk"):

            os.remove(self.data_dir+"missing_ratio-0.2_seed-1.pk")

        # Create folder

        current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

        # define these as variables

        nfold = 1

        config = "acic2018.yaml"

        current_id = "data_pp"

        device = DEVICE

        seed = 1

        testmissingratio = 0.2

        unconditional = 0

        modelfolder = ""

        nsample = 1

        perform_training = 1

        foldername = self.diffpo_path + "/save/acic_fold" + str(nfold) + "_" + current_time + "/"

        # print("model folder:", foldername)

        os.makedirs(foldername, exist_ok=True)

        current_id = "data_pp"

        # print('Start exe_acic on current_id', current_id)

        # Every loader contains "observed_data", "observed_mask", "gt_mask", "timepoints"

        training_size = 1

        train_loader, valid_loader, _ = get_dataloader(

            seed=seed,

            nfold=nfold,

            batch_size=self.config["train"]["batch_size"],

            missing_ratio=testmissingratio,

            dataset_name = self.config["dataset"]["data_name"],

            current_id = current_id,

            training_size = training_size,

            data_path=self.data_dir,

            x_dim=X.shape[1],

        )

        #=======================First train and fix propnet======================

        # Train a propensitynet on this dataset

        propnet = load_data(dataset_name = self.config["dataset"]["data_name"], current_id=current_id, x_dim=X.shape[1], data_path=self.data_dir)

        # frozen the trained_propnet

        # print('Finish training propnet and fix the parameters')

        propnet.eval()

        # ========================================================================

        propnet = propnet.to(device)

        model = TabCSDI(self.config, self.device).to(self.device)

        # Train the model

        train(

            model,

            self.config["train"],

            train_loader,

            valid_loader=valid_loader,

            valid_epoch_interval=self.config["train"]["valid_epoch_interval"],

            foldername=foldername,

            propnet = propnet

        )

        directory = self.diffpo_path + "/save_model/" + current_id

        if not os.path.exists(directory):

            os.makedirs(directory)

        # # load model

        # model.load_state_dict(torch.load(directory + "/model_weights.pth"))

        # save model

        torch.save(model.state_dict(), directory + "/model_weights.pth")

    # predict function with bool return_po and return potential outcome if true

    def predict(self, X: np.ndarray, T0: np.ndarray = None, T1: np.ndarray = None, outcomes: np.ndarray = None) -> np.ndarray:

        """

        Predict the treatment effect using the DiffPO estimator.

        """

        # Store data for their pipeline

        data_dir = self.data_dir

        data, mask = self.reshape_data(X, T0, outcomes)

        data.to_csv(data_dir+"acic2018_norm_data/data_pp_test.csv", index=False)

        mask.to_csv(data_dir+"acic2018_mask/data_pp_test.csv", index=False)

        # Remove old files

        if os.path.exists(data_dir+"missing_ratio-0.2_seed-1_current_id-data_max-min_norm.pk"):

            os.remove(data_dir+"missing_ratio-0.2_seed-1_current_id-data_max-min_norm.pk")

        if os.path.exists(data_dir+"missing_ratio-0.2_seed-1.pk"):

            os.remove(data_dir+"missing_ratio-0.2_seed-1.pk")

        # Create folder

        current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

        # define these as variables

        nfold = 1

        current_id = "data_pp_test"

        current_id_train = "data_pp"

        seed = 1

        testmissingratio = 0.2

        nsample = 50

        perform_training = 1

        foldername = "./save/acic_fold" + str(nfold) + "_" + current_time + "/"

        # print("model folder:", foldername)

        os.makedirs(foldername, exist_ok=True)

        # Every loader contains "observed_data", "observed_mask", "gt_mask", "timepoints"

        training_size = 0

        _,_,test_loader = get_dataloader(

            seed=seed,

            nfold=nfold,

            batch_size=1,

            missing_ratio=testmissingratio,

            dataset_name = self.config["dataset"]["data_name"],

            current_id = current_id,

            training_size = training_size,

            data_path=data_dir,

            x_dim=X.shape[1],

        )

        # load model

        directory = self.diffpo_path + "/save_model/" + current_id_train

        os.makedirs(directory, exist_ok=True)

        model = TabCSDI(self.config, self.device).to(self.device)

        model.load_state_dict(torch.load(directory + "/model_weights.pth"))

        # get cates

        return self.evaluate(model, test_loader, nsample, foldername=foldername)

    def predict_outcomes(self, X: np.ndarray, T0: np.ndarray = None, T1: np.ndarray = None, outcomes: np.ndarray = None) -> np.ndarray:

        """

        Predict the potential outcomes using the DiffPO estimator.

        """

        # add outer dimension to self.pred_outcomes

        return self.pred_outcomes.cpu().numpy().reshape(self.pred_outcomes.shape[0], self.pred_outcomes.shape[1], 1)

    def explain(self, X: np.ndarray, background_samples: np.ndarray = None, explainer_limit: int = None) -> np.ndarray:

        """

        Explain the treatment effect using the EconML estimator.

        """

        if explainer_limit is None:

            explainer_limit = X.shape[0]

        return self.est.shap_values(X[:explainer_limit], background_samples=None)

    def infer_effect_ci(self, X, T0) -> np.ndarray:

        """

        Infer the confidence interval of the treatment effect using the EconML estimator.

        """

        cates_conf_lbs = self.cate_cis[0]

        cates_conf_ups = self.cate_cis[1]

        temp = cates_conf_lbs[T0 != 0]

        cates_conf_lbs[T0 != 0] = -cates_conf_ups[T0 != 0]

        cates_conf_ups[T0 != 0] = -temp

        return np.array([cates_conf_lbs, cates_conf_ups])

    def evaluate(self, model, test_loader, nsample=100, scaler=1, mean_scaler=0, foldername=""):

        # Control random seed in the current script.

        torch.manual_seed(0)

        np.random.seed(0)

        with torch.no_grad():

            model.eval()

            mse_total = 0

            mae_total = 0

            evalpoints_total = 0

            pehe_test = AverageMeter()

            y0_test = AverageMeter()

            y1_test = AverageMeter()

            # for uncertainty

            y0_samples = []

            y1_samples = []

            y0_true_list = []

            y1_true_list = []

            ite_samples = []

            ite_true_list = []

            pred_ites = []

            pred_y0s = []

            pred_y1s = []

            for batch_no, test_batch in enumerate(test_loader, start=1):

                # Get model outputs

                output = model.evaluate(test_batch, nsample) 

                samples, observed_data, target_mask, observed_mask, observed_tp = output

                # Extract relevant quantities

                y0_samples.append(samples[:,:,0]) 

                y1_samples.append(samples[:,:,1]) 

                ite_samples.append(samples[:,:,1] - samples[:,:,0])

                # Get point estimation through median

                est_data = torch.median(samples, dim=1).values

                # Get true ite

                obs_data = observed_data.squeeze(1)

                true_ite = obs_data[:, 2] - obs_data[:, 1] 

                ite_true_list.append(true_ite)

                # Get predicted ite

                pred_y0 = est_data[:, 0]

                pred_y1 = est_data[:, 1]

                pred_y0s.append(pred_y0)

                pred_y1s.append(pred_y1)

                y0_true_list.append(obs_data[:, 1])

                y1_true_list.append(obs_data[:, 2])

                pred_ite = pred_y1 - pred_y0

                pred_ites.append(pred_ite)

                #y0_test.update(diff_y0, obs_data.size(0))    

                #diff_y0 = np.mean((pred_y0.cpu().numpy()-obs_data[:, 1].cpu().numpy())**2)

                #y1_test.update(diff_y1, obs_data.size(0)) 

                #diff_y1 = np.mean((pred_y1.cpu().numpy()-obs_data[:, 2].cpu().numpy())**2)

                #pehe_test.update(diff_ite, obs_data.size(0))    

                #diff_ite = np.mean((true_ite.cpu().numpy()-est_ite.cpu().numpy())**2)

#---------------uncertainty estimation-------------------------

            pred_samples_y0 = torch.cat(y0_samples, dim=0)

            pred_samples_y1 = torch.cat(y1_samples, dim=0)

            pred_samples_ite = torch.cat(ite_samples, dim=0)

            truth_y0 = torch.cat(y0_true_list, dim=0) 

            truth_y1 = torch.cat(y1_true_list, dim=0) 

            truth_ite = torch.cat(ite_true_list, dim=0)

            prob_0, median_width_0 = self.compute_interval(pred_samples_y0, truth_y0)

            prob_1, median_width_1 = self.compute_interval(pred_samples_y1, truth_y1)

            prob_ite, median_width_ite = self.compute_interval(pred_samples_ite, truth_ite)

            self.cate_cis = torch.zeros(2, pred_samples_ite.shape[0], 1) # confidence intervals, dim: 2, n, dim_Y

            for i in range(pred_samples_ite.shape[0]):

                lower_quantile, upper_quantile, in_quantiles = self.check_intervel(confidence_level=0.95, y_pred= pred_samples_ite[i, :], y_true=truth_ite[i])

                self.cate_cis[0, i, 0] = lower_quantile

                self.cate_cis[1, i, 0] = upper_quantile

    #----------------------------------------------------------------

        pred_ites = torch.cat(pred_ites, dim=0)

        pred_y0s = torch.cat(pred_y0s, dim=0)

        pred_y1s = torch.cat(pred_y1s, dim=0)

        #np.zeros((X.shape[0], self.cfg.simulator.num_T, self.cfg.simulator.dim_Y))

        self.pred_outcomes = torch.cat([pred_y0s.unsqueeze(1), pred_y1s.unsqueeze(1)], dim=1)

        self.cate_cis = self.cate_cis.cpu().numpy()

        return pred_ites

    def check_intervel(self, confidence_level, y_pred, y_true):

        lower = (1 - confidence_level) / 2

        upper = 1 - lower

        lower_quantile = torch.quantile(y_pred, lower)

        upper_quantile = torch.quantile(y_pred, upper)

        in_quantiles = torch.logical_and(y_true >= lower_quantile, y_true <= upper_quantile)

        return lower_quantile, upper_quantile, in_quantiles

    def compute_interval(self, po_samples, y_true):

        counter = 0

        width_list = []

        for i in range(po_samples.shape[0]):

            lower_quantile, upper_quantile, in_quantiles = self.check_intervel(confidence_level=0.95, y_pred= po_samples[i, :], y_true=y_true[i])

            if in_quantiles == True:

                counter+=1

            width = upper_quantile - lower_quantile

            width_list.append(width.unsqueeze(0))

        prob = (counter/po_samples.shape[0])

        all_width = torch.cat(width_list, dim=0)

        median_width = torch.median(all_width, dim=0).values

        return prob, median_width