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)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = t.bfloat16
model = HookedSAEBlipConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = model.to(device, dtype=dtype)

num_workers=16

# hf_dataset="yerevann/coco-karpathy"
# local_train_path="./COCO-Dataset/train"
# local_val_path="./COCO-Dataset/val"


hf_dataset="pixparse/cc3m-wds"
local_train_path="./CC3M-Dataset/cc3m_images/train"
local_val_path="./CC3M-Dataset/cc3m_images/val"

tok_name="Salesforce/blip-image-captioning-base"
batch_size=64
max_length=512
filter_seq_length = 30

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 = tok_name,
)

# Load and process image
image, text_prompt = load_image_and_text_from_folder("COCO-Dataset/filtered_val/hallucinated/475466")
# display(image)
print(text_prompt)
# 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_sae_paths = [
    # "cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.{layer}.attention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt",
    "cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.{layer}.crossattention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt",
]

vision_sae_paths = [
    "cc3m_checkpoints/topk_32.0_32_vision_model.encoder.layers.{layer}.hook_resid_post_0.001_256_0.03125_42.ckpt",
]

text_saes = {} # layer: list[str, sae]
for path in text_sae_paths:
    for layer in range(model.cfg.n_layers):
        if layer not in text_saes:
            text_saes[layer] = []
            
        sae = load_sae_model(
            file_path=path.format(layer=layer),
            model_type="blip",
            hook_type="text",
            device=device,
        ).to(dtype)
        text_saes[layer].append((sae.cfg.hook_name, sae))
        
vision_saes = {} # layer: list[str, sae]
for path in vision_sae_paths:
    for layer in range(model.cfg.n_layers):
        if layer not in vision_saes:
            vision_saes[layer] = []
            
        sae = load_sae_model(
            file_path=path.format(layer=layer),
            model_type="blip",
            hook_type="vision",
            device=device,
        ).to(dtype)
        vision_saes[layer].append((sae.cfg.hook_name, sae))
        
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=True, intra=True)


use_error_term = True
pass_through_grad = True
reverse_pruning = False
error_sae = False
gradient_mode = "standard"
edge_gradient_mode = "gradient"
ig_steps = 10 # IG steps
transfer_grad = True

prune_type = "attrib"
cut_mode = "edge"
scoring_mode = "abs"

threshold_type = "value"
threshold = 1

node_threshold_type = "number"
node_threshold = 250 # only for edge attribution

# 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 = -2
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=[1, 0.5],
    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=None,
    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=None,
)

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("sae_graph_nonhal_2.html")