import warnings
import logging
import json
import sys

sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/")
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/")
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/models")
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/models")

import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd

from TimeLLM import Model as TimeLLMModel
from unitime import UniTime as UniTimeModel

IMPLEMENTED_BASELINES = [TimeLLMModel, UniTimeModel]

from typing import Optional, Union, Dict, Callable, Iterable

def truncate_mse_loss(future_time, future_pred):
    # Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
    min_length = min(future_time.shape[-1], future_pred.shape[-1])
    return F.mse_loss(future_time[...,:min_length], future_pred[...,:min_length])

def truncate_mae_loss(future_time, future_pred):
    # Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
    min_length = min(future_time.shape[-1], future_pred.shape[-1])
    return F.l1_loss(future_time[...,:min_length], future_pred[...,:min_length])

class DotDict(dict):
    """dot.notation access to dictionary attributes"""
    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

def find_pred_len_from_path(path: str) -> int:
    if "pl_96" or "pl96" in path: pred_len = 96
    elif "pl_192" or "pl192" in path: pred_len = 192
    elif "pl_336" or "pl336" in path: pred_len = 336
    elif "pl720" or "pl720" in path: pred_lent = 720
    else:
        raise ValueError(f"Could not determine prediction length of model from path {path}. Expected path to contain a substring of the form 'pl_{{pred_len}}' or 'pl{{pred_len}}'.")

    return pred_len

def find_model_name_from_path(path: str) -> str:
    path = path.lower()
    if "time-llm" in path or "timellm" in path: model_name = "time-llm"
    elif "unitime" in path: model_name = "unitime"
    else:
        raise ValueError(f"Could not determine model name from path {path}. Expected path to contain either 'time-llm', 'timellm', or 'unitime'.")

    return model_name

TIME_LLM_CONFIGS = DotDict({
    "task_name" : "long_term_forecast", "seq_len" : 512, "enc_in" : 7, "d_model" : 32, "d_ff" : 128, "llm_layers" : 32, "llm_dim" : 4096, 
    "patch_len" : 16, "stride" : 8, "llm_model" : "LLAMA", "llm_layers" : 32, "prompt_domain" : 1, "content" : None, "dropout" : 0.1, 
    "d_model" : 32, "n_heads" : 8, "enc_in" : 7
    })

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
UNITIME_CONFIGS = DotDict({
    "max_token_num" : 17, "mask_rate" : 0.5, "patch_len" : 16, "max_backcast_len" : 96, "max_forecast_len" : 720, "logger" : logger,
    "model_path" : "gpt2", "lm_layer_num" : 6, "lm_ft_type" : "freeze", "ts_embed_dropout" : 0.3, "dec_trans_layer_num" : 2, "dec_head_dropout" : 0.1,
})

class TimeLLMStarCasterWrapper(nn.Module):

    def __init__(self, time_llm_model):
        super().__init__()

        assert isinstance(time_llm_model, TimeLLMModel), f"TimeLLMStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(time_llm_model)}"
        self.base_model = time_llm_model
    
    def forward(self, past_time, context):
        self.base_model.description = context
        return self.base_model(x_enc=past_time.unsqueeze(-1), x_mark_enc=None, x_dec=None, x_mark_dec=None).squeeze(-1)
    
class UniTimeStarCasterWrapper(nn.Module):

    def __init__(self, unitime_model):
        super().__init__()

        assert isinstance(unitime_model, UniTimeModel), f"UniTimeStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(unitime_model)}"
        self.base_model = unitime_model

    def forward(self, past_time, context):
        past_time = past_time.unsqueeze(-1)
        mask = torch.ones_like(past_time)
        data_id = -1
        seq_len = 96
        stride = 16

        info = (data_id, seq_len, stride, context[:17])
        return self.base_model(info=info, x_inp=past_time, mask=mask).squeeze(-1)

class StarCasterBaseline(nn.Module):

    def __init__(self, model):
        super().__init__()

        # TODO: Make this more extendable
        if type(model) not in IMPLEMENTED_BASELINES:
            raise NotImplementedError(f"StarCasterBaseline currently only handles models of type {IMPLEMENTED_BASELINES}.")
        
        self.base_model = model
        if isinstance(self.base_model, TimeLLMModel):
            self.wrapped_model = TimeLLMStarCasterWrapper(self.base_model)
        if isinstance(self.base_model, UniTimeModel):
            self.wrapped_model = UniTimeStarCasterWrapper(self.base_model)
    
    def forward(self, past_time, context):
        return self.wrapped_model(past_time, context)
    
    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        return self.base_model.load_state_dict(state_dict, strict, assign)


class EvaluationPipeline:

    def __init__(
        self, 
        dataset: Iterable,
        model: TimeLLMModel, 
        metrics: Optional[Union[Callable, Dict[str, Callable]]] = None
    ):
        self.dataset = dataset
        self.metrics = metrics if metrics is not None else {"mse_loss" : truncate_mse_loss}

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        if self.device == "cpu":
            warnings.warn("Warning: No CUDA device detected, proceeding with EvaluationPipeline on CPU .....")

        self.model = StarCasterBaseline(model).to(self.device)
    
    # TODO: This method needs to be replaced to handle actual StarCaster benchmark
    def get_evaluation_loader(self) -> Iterable:
        samples = []
        for sample in self.dataset.values():
            past_time = torch.from_numpy(sample["past_time"].to_numpy().T).float().to(self.device)
            future_time = torch.from_numpy(sample["future_time"].to_numpy().T).float().to(self.device)
            context = sample["context"]

            samples.append([past_time, future_time, context])

        return samples

    def compute_loss(self, future_time, future_pred):
        return {m_name : m(future_time, future_pred) for m_name, m in self.metrics.items()}
    
    def evaluation_step(self, past_time, future_time, context):
        with torch.no_grad():
            future_pred = self.model(past_time, context)
            loss = self.compute_loss(future_time, future_pred)
        return loss, future_pred
            
    @torch.no_grad()
    def eval(self):
        model.eval()
        infer_dataloader = self.get_evaluation_loader()
        losses, predictions = {m_name : [] for m_name in self.metrics.keys()}, []
        for past_time, future_time, context in infer_dataloader:
            loss_dict, preds = self.evaluation_step(past_time, future_time, context)

            for m_name, loss in loss_dict.items(): losses[m_name].append(loss)
            predictions.append(preds)

        model.train()
        return losses, predictions

if __name__ == "__main__":
    # from argparse import ArgumentParser

    # parser = ArgumentParser()

    # parser.add_argument("--data_path", type=str, required=True)
    # parser.add_argument("--ckpt_path", type=str, default=None)
    
    # args = parser.parse_args()

    # args = TIME_LLM_CONFIGS
    args = DotDict(dict())

    # args.ckpt_path = "./Time-LLM/checkpoints/long_term_forecast_ETTh1_512_96_TimeLLM_ETTh1_ftM_sl512_ll48_pl96_dm32_nh8_el2_dl1_df128_fc3_ebtimeF_Exp_0-TimeLLM-ETTh1/best_checkpoint/pytorch_model/mp_rank_00_model_states.pt"
    args.ckpt_path = "/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/outputs/checkpoint_gpt2-small_full_etth1-96_instruct_6_2_0.5_96/model_s2036.pth"
    args.data_path = "./example_data_dict_simple_dtypes.pkl"

    dataset = pd.read_pickle(args.data_path)
    # args.pred_len = find_pred_len_from_path(args.ckpt_path)
    # args.model_name = find_model_name_from_path(args.ckpt_path)
    args.pred_len = 96
    args.model_name = "unitime" # "time-llm"
    
    if args.model_name == "time-llm":
        args.update(TIME_LLM_CONFIGS)
    elif args.model_name == "unitime":
        args.update(UNITIME_CONFIGS)

    print(f"Initializing model from config:\n{args} .....")

    if args.model_name == "time-llm":
        model = TimeLLMModel(args)
    elif args.model_name == "unitime":
        model = UniTimeModel(args)

    if args.ckpt_path is not None:
        print(f"Loading model checkpoint from path {args.ckpt_path} .....")
        ckpt = torch.load(args.ckpt_path)
        if args.model_name == "time-llm":
            model.load_state_dict(ckpt["module"]) # TODO: Change this to not be specific to the Time-LLM checkpoint
        elif args.model_name == "unitime":
            model.load_state_dict(ckpt)

    pipeline = EvaluationPipeline(dataset, model, metrics={"mse_loss" : truncate_mse_loss, "mae_loss" : truncate_mae_loss})

    print(f"Evaluating .....")
    losses, predictions = pipeline.eval()
    print(f"Got losses: {losses}")
    print(f"Predictions has shape: {[pred.shape for pred in predictions]}")