import os
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import csv
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.gpt2.modeling_gpt2 import GPT2MLP
from transformers.models.gptj.modeling_gptj import GPTJMLP
from transformers.models.llama.modeling_llama import LlamaMLP

from rome.layer_stats import layer_stats
from util import nethook
from util.generate import generate_fast
from util.globals import *

from .compute_ks import compute_ks, compute_seq_ks
from .compute_z import compute_z, get_module_input_output_at_words, find_fact_lookup_idx
from .rewrited_mlp import CustomGPT2MLP, CustomGPTJMLP, CustomLlamaMLP

from .NDEdit_hparams import NDEditHyperParams
# Cache variable(s)
CONTEXT_TEMPLATES_CACHE = None
COV_CACHE = {}
N_CACHE = {}

def apply_rewrited_mlp(old_mlp, config):
    if isinstance(old_mlp, GPT2MLP):
        return CustomGPT2MLP(old_mlp, config)
    elif isinstance(old_mlp, GPTJMLP):
        return CustomGPTJMLP(old_mlp, config)
    elif isinstance(old_mlp, LlamaMLP):
        return CustomLlamaMLP(old_mlp, config) 
    else:
        print('Not supporting this model')
        assert True




def apply_NDEdit_to_model(
    model: AutoModelForCausalLM,
    tok: AutoTokenizer,
    requests: List[Dict],
    hparams: NDEditHyperParams,
    cache_template: Optional[str] = None,
    cache_c = None,
    N = None,
) -> Dict[str, Tuple[torch.Tensor]]:
    """
    Executes the MEMIT update algorithm for the specified update at the specified layer
    Invariant: model at beginning of function == model at end of function
    """

    # Update target and print info
    requests = deepcopy(requests)
    for i, request in enumerate(requests):
        if request["target_new"]["str"][0] != " ":
            # Space required for correct tokenization
            requests[i]["target_new"]["str"] = " " + request["target_new"]["str"]
    for request in requests[:10]:
        print(
            f"MEMIT request sample: "
            f"[{request['prompt'].format(request['subject'])}] -> [{request['target_new']['str']}]"
        )

    # Retrieve weights that user desires to change
    weights = {
        f"{hparams.rewrite_module_tmp.format(layer)}.weight": nethook.get_parameter(
            model, f"{hparams.rewrite_module_tmp.format(layer)}.weight"
        )
        for layer in hparams.layers
    }
    weights_copy = {k: v.detach().clone() for k, v in weights.items()}
    # Compute z for final layer
    context_templates = get_context_templates(model, tok)
    z_layer = hparams.layers[-1]
    z_list = []

    for request in requests:
        # Retrieve k/v pair if already stored in cache
        cache_fname = (
            Path(
                str(cache_template).format(
                    z_layer, hparams.clamp_norm_factor, request["case_id"]
                )
            )
            if cache_template is not None
            else None
        )
        data_loaded = False
        if (
            cache_fname is not None  # Require cache template
            and cache_fname.exists()  # Cache file must exist
        ):
            try:
                data = np.load(cache_fname)
                z_list.append(torch.from_numpy(data["v_star"]).to("cuda"))
                data_loaded = True
            except Exception as e:
                print(f"Error reading cache file due to {e}. Recomputing...")

        # Compute k/v pair if not loaded from cache
        if not data_loaded:
            cur_z = compute_z(
                model,
                tok,
                request,
                hparams,
                z_layer,
                context_templates,
            )

            z_list.append(cur_z)

            if cache_fname is not None:
                cache_fname.parent.mkdir(exist_ok=True, parents=True)
                np.savez(
                    cache_fname,
                    **{
                        "v_star": cur_z.detach().cpu().numpy(),
                    },
                )
                print(f"Cached k/v pair at {cache_fname}")
    zs = torch.stack(z_list, dim=1).detach().T

    for i, layer in enumerate(hparams.layers):
        print(f"\n\nLAYER {layer}\n")

        # Get current model activations
        layer_ks_list, cur_zs_list = [], []
        for i in range(0, len(requests), hparams.batch_size):
            batch_requests = requests[i:i + hparams.batch_size]
            batch_layer_ks = compute_ks(model, tok, batch_requests, hparams, layer, context_templates).detach()
            batch_cur_zs = get_module_input_output_at_words(
                model,
                tok,
                z_layer,
                context_templates=[request["prompt"] for request in batch_requests],
                words=[request["subject"] for request in batch_requests],
                module_template=hparams.layer_module_tmp,
                fact_token_strategy=hparams.fact_token,
            )[1].detach()
            layer_ks_list.append(batch_layer_ks)
            cur_zs_list.append(batch_cur_zs)
        layer_ks = torch.cat(layer_ks_list, dim=0)
        cur_zs = torch.cat(cur_zs_list, dim=0)
        target_v = zs - cur_zs 

        print(f"Writing {layer_ks.size(0)} key/value pair(s) into layer {layer}")

        # original_mlp = model.transformer.h[layer].mlp
        cur_mlp = nethook.get_module(model, hparams.mlp_module_tmp.format(layer))
        assert layer_ks.shape[0] == target_v.shape[0]
        try:
            cur_mlp.update_KV(layer_ks, target_v)
        except:
            try:
                new_mlp =  apply_rewrited_mlp(cur_mlp, model.config)
            except:
                new_mlp =  apply_rewrited_mlp(cur_mlp, model.transformer.config)
            new_mlp.update_KV(layer_ks, target_v)
            cur_layer_module = nethook.get_module(model, hparams.layer_module_tmp.format(layer))
            setattr(cur_layer_module, 'mlp', new_mlp)

        # Adjust update matrix shape

        del layer_ks,  z_list, layer_ks_list, cur_zs_list, cur_zs, target_v
        torch.cuda.empty_cache()

    print(f"Deltas successfully computed for {list(weights.keys())}")
    return model


def get_cov(
    model: AutoModelForCausalLM,
    tok: AutoTokenizer,
    layer_name: str,
    mom2_dataset: str,
    mom2_n_samples: str,
    mom2_dtype: str,
    inv: bool = False,
    force_recompute: bool = False,
) -> torch.Tensor:
    """
    Retrieves covariance statistics, then computes the algebraic inverse.
    Caches result for future use.
    """

    model_name = model.config._name_or_path.replace("/", "_")
    key = (model_name, layer_name)

    print(f"Retrieving covariance statistics for {model_name} @ {layer_name}.")
    if key not in COV_CACHE or force_recompute:
        stat = layer_stats(
            model,
            tok,
            layer_name,
            STATS_DIR,
            mom2_dataset,
            to_collect=["mom2"],
            sample_size=mom2_n_samples,
            precision=mom2_dtype,
            force_recompute=force_recompute,
        )
        COV_CACHE[key] = stat.mom2.moment().float().to("cpu")

    return (
        torch.inverse(COV_CACHE[key].to("cuda")) if inv else COV_CACHE[key].to("cuda")
    )


def upd_matrix_match_shape(matrix: torch.Tensor, shape: torch.Size) -> torch.Tensor:
    """
    GPT-2 and GPT-J have transposed weight representations.
    Returns a matrix that matches the desired shape, else raises a ValueError
    """

    if matrix.shape == shape:
        return matrix
    elif matrix.T.shape == shape:
        return matrix.T
    else:
        raise ValueError(
            "Update matrix computed by MEMIT does not match original weight shape. "
            "Check for bugs in the code?"
        )


def get_context_templates(model, tok):
    global CONTEXT_TEMPLATES_CACHE

    if CONTEXT_TEMPLATES_CACHE is None:
        CONTEXT_TEMPLATES_CACHE = [["{}"]] + [
            [
                f.replace("{", " ").replace("}", " ") + ". {}"
                for f in generate_fast(
                    model,
                    tok,
                    ["The", "Therefore", "Because", "I", "You"],
                    n_gen_per_prompt=n_gen // 5,
                    max_out_len=length,
                )
            ]
            for length, n_gen in [(10, 5)]  # Be careful about changing this.
        ]
        print(f"Cached context templates {CONTEXT_TEMPLATES_CACHE}")

    return CONTEXT_TEMPLATES_CACHE
