# %%
from sae.SAE_Tools import *
import os
from pathlib import Path
from typing import List, Optional
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
from model.blip.hooked_blip import HookedSAEBlipConditionalGeneration
from transformers import BlipProcessor
from huggingface_hub import login
from sae.SAE_Blip_Explaining_Utils import *
from sae.SAE_Trainer import DataConfig
from hallucination.extra_materials.graph_visualizer_blip import *
from hallucination.extra_materials.circuit_utils import *

load_dotenv()   
hf_key = os.getenv('HUGGING_FACE_API_KEY')
login(hf_key)

parser = argparse.ArgumentParser(description="Mechanistic Faithfulness Metric.")
parser.add_argument('--model_name', type=str, default="Salesforce/blip-image-captioning-base", help="The model name.")
parser.add_argument('--tok_name', type=str, default="Salesforce/blip-image-captioning-base", help="The tokenizer name.")
parser.add_argument('--data_path', type=str, default="./COCO-Dataset/filtered_val/hallucinated/52982", help="The path to the data folder containing image and caption.txt.")
parser.add_argument('--target_position', type=int, default=-2, help="The position of the target token to construct the circuit.")
parser.add_argument('--batch_size', type=int, default=16, help="The processing batchsize.")
parser.add_argument('--text_sae_types', type=str, nargs="+", default=["attn"], help="The type of SAE in the text circuit.")
parser.add_argument('--vis_sae_types', type=str, nargs="+", default=["pre"], help="The type of SAE in the vision circuit.")
parser.add_argument('--text_batch', type=int, default=200, help="The text processing batch.")
parser.add_argument('--vis_batch', type=int, default=200, help="The vision processing batch.")
parser.add_argument('--crop_ratio', type=float, nargs="+", default=[1.0], help="The crop sizes to create vision multi-crop dataset.")
parser.add_argument('--inter', type=str_to_bool, default=True, help="Whether to set inter-connection in the circuit.")
parser.add_argument('--intra', type=str_to_bool, default=False, help="Whether to set intra-connection in the circuit.")
parser.add_argument('--gradient_mode', type=str, default="standard", help="The type of gradient to compute node score")
parser.add_argument('--edge_gradient_mode', type=str, default="gradient", help="The type of gradient to compute edge score")
parser.add_argument('--num_nodes', type=int, default=100, help="The number of nodes in the circuit.")
parser.add_argument('--scoring_mode', type=str, default="less", help="The number of nodes in the circuit.")
parser.add_argument('--ig_steps', type=int, default=10, help="The number of steps for integrated gradients.")
parser.add_argument('--filter_seq_length', type=int, default=30, help="Threshold to filter long sequence caption.")
parser.add_argument('--save_name', type=str, default="sae_graph", help="Graph save name.")
parser.add_argument('--dtype', type=str, default="bfloat16", help="The dtype.")
parser.add_argument('--device_id', type=int, default=0, help="The id of GPU.")

args = parser.parse_args()

device = f'cuda:{args.device_id}' if torch.cuda.is_available() else 'cpu'
dtype = str_to_dtype(args.dtype)
model = HookedSAEBlipConditionalGeneration.from_pretrained(args.model_name)
processor = BlipProcessor.from_pretrained(args.tok_name)
model = model.to(device, dtype=dtype)

num_workers=4
hf_dataset="yerevann/coco-karpathy"
local_train_path="./COCO-Dataset/train_rest"
local_val_path="./COCO-Dataset/val"
batch_size=args.batch_size
max_length=512
filter_seq_length = args.filter_seq_length

data_config = DataConfig(
    batch_size=batch_size,
    hf_dataset=hf_dataset,
    local_train_path=local_train_path,
    local_val_path=local_val_path,
    num_workers=num_workers,
    max_length=max_length, # the processor of blip only allow max tokens (fixed)
    processor=args.tok_name,
)

# Load and process image
image, text_prompt = load_image_and_text_from_folder(args.data_path)

# Prepare inputs and move to device
inputs = processor(images=image, text=text_prompt, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()} 

# %%
text_saes, vision_saes = load_saes(
    model,
    model_type="blip",
    text_sae_types=args.text_sae_types,
    vis_sae_types=args.vis_sae_types,
    device=device,
    dtype=dtype,
)

# %%
from hallucination.extra_materials.graph.Feature_Graph_Blip import Feature_Graph_Blip
from hallucination.extra_materials.Utils import Pruner

fg = Feature_Graph_Blip(
    model,
    text_saes,
    vision_saes,
    use_error_term=True
)
_ = fg.build_default_connection(seq_length=inputs['input_ids'].shape[1], token_wise=True, inter=args.inter, intra=args.intra)

# %%
pass_through_grad = True
reverse_pruning = False
error_sae = False
gradient_mode = args.gradient_mode
edge_gradient_mode = args.edge_gradient_mode
ig_steps = args.ig_steps
transfer_grad = True

prune_type = "attrib"
cut_mode = "edge"
scoring_mode = args.scoring_mode

threshold_type = "value"
threshold = 2

node_threshold_type = "number"
node_threshold = args.num_nodes

# random input
corrupt_inputs = {
    key: t.zeros_like(value) for key, value in inputs.items()
}
corrupt_logits, corrupt_cache = fg.run_model(corrupt_inputs)

pos = args.target_position
print("Targeting token:", processor.tokenizer.decode(inputs['input_ids'][0, pos]))

metric = lambda logits, token=inputs['input_ids'][0, pos]: logits[:, pos-1, token]
    
# fwd_cache, bwd_cache = fg.forward_backward_gradient(inputs, corrupt_cache, metric)

pruner = Pruner(fg, metric, device, verbose=False)

# prune node
pruner( # type: ignore
    inputs, 
    corrupt_inputs, 
    node_threshold, 
    prune_type=prune_type, 
    cut_mode="node", 
    scoring_mode=scoring_mode, 
    threshold_type=node_threshold_type,
    modify_inplace=True,   # modify inplace
    return_type="retained",
    reverse_pruning=reverse_pruning,
    verbose=False,
    pass_through_grad=pass_through_grad,
    gradient_mode=gradient_mode,
    transfer_grad=transfer_grad,
    steps=ig_steps,
)

# prune edge
pruned_graph, retained_components = pruner( # type: ignore
    inputs, 
    corrupt_inputs, 
    threshold, 
    prune_type=prune_type, 
    cut_mode=cut_mode, 
    scoring_mode=scoring_mode, 
    threshold_type=threshold_type,
    modify_inplace=True,   # modify inplace
    return_type="retained",
    reverse_pruning=reverse_pruning,
    verbose=True,
    pass_through_grad=pass_through_grad,
    gradient_mode=gradient_mode,
    edge_gradient_mode=edge_gradient_mode,
)

for key, val in fg.nodes.items():
    fg.nodes[key] = val.to("cpu")
for key, val in fg.node_scores.items():
    fg.node_scores[key] = val.to("cpu")
for key, val in fg.edges.items():
    for key2, val2 in val.items():
        val[key2] = val2.to("cpu")
for key, val in fg.edge_scores.items():
    for key2, val2 in val.items():
        val[key2] = val2.to("cpu")
        
t.cuda.empty_cache()

# %%
dict_acts = {}
# VISION
vision_sae_list = []
for saes in vision_saes.values():
    vision_sae_list.extend([sae[1] for sae in saes])    

nocap_train, nocap_val = load_lvlm_data_nocap(
    config=data_config
)

processed_ds = DebatchNoCapDataset(nocap_val, processor=data_config.processor)

multi_crop_dataset = MultiScaleCropDataset(
    original_dataset=processed_ds,
    img_size=model.config.vision_config.image_size,
    crop_ratios=args.crop_ratio,
    stride_ratio=0.5,
    resize_to=model.config.vision_config.image_size,
)

dataloader = DataLoader(multi_crop_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

cache_dict, _ = cache_vision_sae_lvlm(
    vision_sae_list,
    model,
    dataloader,
    device,
    filter_seq_length=filter_seq_length,
    stop_at_batch=args.vis_batch,
    return_toks=False,
)

dict_acts = dict_acts | cache_dict

# TEXT
text_sae_list = []
for saes in text_saes.values():
    text_sae_list.extend([sae[1] for sae in saes])  
    
train_loader, val_loader = load_lvlm_data(data_config)

cache_dict, data_toks = cache_sae_lvlm(
    text_sae_list,
    model,
    val_loader,
    device,
    filter_seq_length=filter_seq_length,
    stop_at_batch=args.text_batch,
)

dict_acts = dict_acts | cache_dict

del train_loader, val_loader, dataloader, nocap_train, nocap_val

# %%
graph_vis = SAEGraphVisualizer(
    graph_obj=fg, 
    model=model,
    processor=processor, 
    sae_dict=fg.dict_saes,
    tokens=inputs['input_ids'][0],
    full_data_dict=dict_acts,
    data_toks=data_toks,
    dataset=multi_crop_dataset
)

html_code = graph_vis.generate_html(args.save_name)