import math
from dataclasses import dataclass


import torch
from datasets import IterableDataset, load_dataset, concatenate_datasets
from datasets.formatting.formatting import LazyBatch
from jaxtyping import Int
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, PreTrainedTokenizerBase
from typing import Callable, List
from .SAE_Trainer import DataConfig

from PIL import Image
import os
from pathlib import Path
from transformers import BlipProcessor, LlavaProcessor
import torch.nn.functional as F
from tqdm import tqdm

# Create COCO Datasets

def find_local_image(split: str, config: DataConfig, id: int | str) -> str | None:
    if split == 'validation':
        root_path = Path(config.local_val_path)
    else:
        root_path = Path(config.local_train_path)

    img_filename = f"{str(id)}.jpg"
    p = root_path / img_filename
    if p.exists():
        return str(p)
    return None


class COCOLocalDataset(Dataset):
    """
    Dataset for COCO-style data with multiple captions per image.
    Expands the dataset by creating one sample per caption (repeating the image).
    Loads images from local filesystem using the config and split.
    """
    def __init__(self,
                 hf_dataset,
                 id_field: str,
                 txt_field: str,
                 split: str,
                 config: DataConfig
                 ):
        self.hf = hf_dataset
        self.id_field = id_field
        self.txt_field = txt_field
        self.split = split
        self.config = config

        self.ds_index = []
        for row_idx in range(len(self.hf)):
            row = self.hf[row_idx]
            sentences = row.get(self.txt_field, [])
            for cap_idx in range(len(sentences)):
                self.ds_index.append((row_idx, cap_idx))

    def __len__(self):
        return len(self.ds_index)

    def _get_id(self, row):
        return row[self.id_field]  

    def __getitem__(self, idx):
        row_idx, cap_idx = self.ds_index[idx]
        row = self.hf[row_idx]
        cocoid = self._get_id(row)
        img_path = find_local_image(
            split=self.split,
            config=self.config,
            id=cocoid)
        if img_path is None:
            raise FileNotFoundError(f"Image not found for ImgID {cocoid}")
        img = Image.open(img_path).convert("RGB")
        caption = row[self.txt_field][cap_idx]

        return {
            "image": img,
            "imgid": cocoid,
            "caption": caption,
        }
        
class CC3MLocalDataset(Dataset):
    """
    CC3M dataset: one caption per image (txt), so no ds_index expansion.
    O(1) initialization: does not iterate through the HF dataset in __init__.
    """
    def __init__(self, 
                 hf_dataset, 
                 id_field: str, 
                 txt_field: str, 
                 split: str, 
                 config):
        self.hf = hf_dataset
        self.id_field = id_field      
        self.txt_field = txt_field    
        self.split = split
        self.config = config

    def __len__(self):
        return len(self.hf)

    def _get_id(self, row):
        return row[self.id_field]

    def __getitem__(self, idx):
        row = self.hf[idx]
        id_ = self._get_id(row)

        img_path = find_local_image(split=self.split, 
                                    config=self.config, 
                                    id=id_)
        if img_path is None:
            raise FileNotFoundError(f"Image not found for ImgID {id_}")

        img = Image.open(img_path).convert("RGB")

        caption = row[self.txt_field]

        return {
            "image": img,
            "imgid": id_,
            "caption": caption,
        }


def coco_dataset(config: DataConfig, split: str):
    if split == 'train':
        split = 'train'
    elif split == 'validation':
        split = 'validation'
    elif split == 'trainrest':
        split = 'train+restval'
        
    hf_ds = load_dataset(config.hf_dataset, split=split)
    
    id_field = "cocoid"
    txt_field = "sentences"
    
    return COCOLocalDataset(
        hf_dataset=hf_ds,
        id_field=id_field,
        txt_field=txt_field,
        split=split,
        config=config,
    )

def cc3m_dataset(config: DataConfig, split: str):
    if split == 'train':
        split = 'train'
    elif split == 'validation':
        split = 'validation'
        
    hf_ds = load_dataset(config.hf_dataset, split=split)
    id_field = "__key__"
    txt_field = "txt"
    
    return CC3MLocalDataset(
        hf_dataset=hf_ds,
        id_field=id_field,
        txt_field=txt_field,
        split=split,
        config=config,
    )

def load_lvlm_data(
    config: DataConfig,
) -> tuple[DataLoader, DataLoader]:

    if "llava" in config.processor:
        processor = LlavaProcessor.from_pretrained(config.processor)
        def format_text(caption):
            return f"USER: <image>\nDescribe this image. \nASSISTANT: {caption}"
    else:
        processor = BlipProcessor.from_pretrained(config.processor)
        def format_text(caption):
            return caption

    def collate_fn_lvlm(batch: List[dict]):
        images = [b["image"] for b in batch]
        caption = [format_text(b["caption"]) for b in batch]
        processed = processor(images=images, text=caption, return_tensors="pt", padding=True)

        out = {
            **processed,
            "imgids": [b["imgid"] for b in batch],
            "captions": [b["caption"] for b in batch],
        }
        return out
    

    if 'coco' in config.hf_dataset:
        print("Loading COCO dataset...")
        train_ds = coco_dataset(config, "train")
        val_ds = coco_dataset(config, "validation")
        
    elif 'cc3m' in config.hf_dataset:
        print("Loading CC3M dataset...")
        train_ds = cc3m_dataset(config, "train")
        val_ds = cc3m_dataset(config, "validation")
        
    train_dataloader = DataLoader(
        train_ds,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=config.num_workers,
        collate_fn=collate_fn_lvlm,
    )

    val_dataloader = DataLoader(
        val_ds,
        batch_size=config.batch_size,
        shuffle=False,
        num_workers=config.num_workers,
        collate_fn=collate_fn_lvlm,   
    )

    return train_dataloader, val_dataloader



def load_multilayer_sae_data(
    config: DataConfig,
) -> tuple[DataLoader, DataLoader]:

    if "llava" in config.processor:
        processor = LlavaProcessor.from_pretrained(config.processor)
        prompt = "USER: <image>\nDescribe this image. \nASSISTANT: "
    else:
        processor = BlipProcessor.from_pretrained(config.processor)
        prompt = f"Describe this image: "
    
    def collate_fn_lvlm(batch: List[dict]):
        images = [b["image"] for b in batch]
        processed = processor(images=images, text=prompt, return_tensors="pt", padding=True)

        out = {
            **processed,
            "imgids": [b["imgid"] for b in batch],
        }
        return out
    

    if 'coco' in config.hf_dataset:
        print("Loading COCO dataset...")
        train_ds = coco_dataset(config, "train")
        val_ds = coco_dataset(config, "validation")
        
    elif 'cc3m' in config.hf_dataset:
        print("Loading CC3M dataset...")
        train_ds = cc3m_dataset(config, "train")
        val_ds = cc3m_dataset(config, "validation")
        
    train_dataloader = DataLoader(
        train_ds,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=config.num_workers,
        collate_fn=collate_fn_lvlm,
    )

    val_dataloader = DataLoader(
        val_ds,
        batch_size=config.batch_size,
        shuffle=False,
        num_workers=config.num_workers,
        collate_fn=collate_fn_lvlm,   
    )

    return train_dataloader, val_dataloader


# Datasets for visualize vision SAE(s)

class LocalDatasetNoCap(Dataset):
    """
        Dataset that maps HuggingFace dataset entries to local images using IDs.
        For each entry, it retrieves the image from the local filesystem based on the ID
    """
    def __init__(self,
                 hf_dataset,
                 id_field: str,
                 split: str,
                 config: DataConfig
                 ):
        self.hf = hf_dataset
        self.id_field = id_field
        self.split = split
        self.config = config

    def __len__(self):
        return len(self.hf)

    def _get_id(self, row):
        return row[self.id_field]  

    def __getitem__(self, idx):
        row = self.hf[idx]
        cocoid = self._get_id(row)
        img_path = find_local_image(
            split=self.split,
            config=self.config,
            id=cocoid)
        if img_path is None:
            raise FileNotFoundError(f"Image not found for ImgID {cocoid}")
        img = Image.open(img_path).convert("RGB")

        return {
            "image": img,
            "imgid": cocoid,
        }


def coco_dataset_nocap(config: DataConfig, split: str):
    if split == 'train':
        split = 'train'
    elif split == 'validation':
        split = 'validation'
    elif split == 'trainrest':
        split = 'train+restval'
        
    hf_ds = load_dataset(config.hf_dataset, split=split)
    id_field = "cocoid"
    

    return LocalDatasetNoCap(
        hf_dataset=hf_ds,
        id_field=id_field,
        split=split,
        config=config,
    )


def cc3m_dataset_nocap(config: DataConfig, split: str):
    if split == 'train':
        split = 'train'
    elif split == 'validation':
        split = 'validation'
        
    hf_ds = load_dataset(config.hf_dataset, split=split)
    id_field = "__key__"
    

    return LocalDatasetNoCap(
        hf_dataset=hf_ds,
        id_field=id_field,
        split=split,
        config=config,
    )
    
def load_lvlm_data_nocap(
    config: DataConfig,
) -> tuple[DataLoader, DataLoader]:
    
    if 'coco' in config.hf_dataset:
        print("Loading COCO nocap dataset...")
        train_ds = coco_dataset_nocap(config, "train")
        val_ds = coco_dataset_nocap(config, "validation")
    
    elif 'cc3m' in config.hf_dataset:
        print("Loading CC3M nocap dataset...")
        train_ds = cc3m_dataset_nocap(config, "train")
        val_ds = cc3m_dataset_nocap(config, "validation")
    
    return train_ds, val_ds

class DebatchNoCapDataset(Dataset):
    """no-caption dataset: processes single images on-the-fly with empty text prompt."""
    
    def __init__(self, dataset, processor):
        self.dataset = dataset
        if "blip" in processor:
            print("Using Blip processor")
            self.processor = BlipProcessor.from_pretrained(processor)
            self.prompt = ""
        elif "llava" in processor:
            print("Using Llava processor")
            self.processor = LlavaProcessor.from_pretrained(processor)
            self.prompt = "USER: <image>\nASSISTANT:"
        else:
            raise ValueError(f"Unsupported processor: {processor}")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        raw_item = self.dataset[idx]
        image = raw_item["image"]

        # Process exactly like batched version but with B=1 and text=""
        processed = self.processor(
            images=image,
            text=self.prompt, 
            return_tensors="pt",
            padding=True,
        )

        return {
            "pixel_values": processed["pixel_values"],
            "input_ids": processed["input_ids"],
            "attention_mask": processed["attention_mask"],
            "imgid": raw_item["imgid"],
        }

class MultiScaleCropDataset(Dataset):
    """Crop images from debatch lvlm datasets to different scales and resize.
    
     Args:
        original_dataset: Debatched dataset.
        img_size: Square image size.
        crop_ratios: List of crop ratios to apply.
        stride_ratio: Stride as a ratio of crop size.
        resize_to: Resize cropped images to square size.
    """
    
    def __init__(
        self,
        original_dataset,
        img_size: int = 384,
        crop_ratios=[1.0, 0.5, 0.25, 0.125],
        stride_ratio: float = 0.5,
        resize_to: int = 384,
    ):
        self.original_dataset = original_dataset
        self.img_size = img_size
        self.resize_to = resize_to
        self.stride_ratio = stride_ratio

        self.crop_configs = []
        for ratio in crop_ratios:
            crop_size = int(img_size * ratio)
            if crop_size == 0:
                continue
            stride = max(int(crop_size * stride_ratio), 1)

            steps_h = max((img_size - crop_size) // stride + 1, 1)
            steps_w = max((img_size - crop_size) // stride + 1, 1)

            for h in range(steps_h):
                for w in range(steps_w):
                    top = h * stride
                    left = w * stride
                    self.crop_configs.append({
                        "crop_size": crop_size,
                        "top": top,
                        "left": left,
                        "ratio": ratio,
                    })

        print(f"Total crops per image: {len(self.crop_configs)}")

        # Total length = len(original) * num_crops_per_image
        self.total_crops = len(original_dataset) * len(self.crop_configs)

    def __len__(self):
        return self.total_crops

    def __getitem__(self, idx):
        # Map global index back to (original_idx, crop_config_idx)
        orig_idx = idx // len(self.crop_configs)
        crop_idx = idx % len(self.crop_configs)
        config = self.crop_configs[crop_idx]

        batch = self.original_dataset[orig_idx]  

        pixel_values = batch["pixel_values"]
        crop = pixel_values[
            :,
            config["top"]:config["top"] + config["crop_size"],
            config["left"]:config["left"] + config["crop_size"]
        ]

        # Resize to target size
        crop = F.interpolate(
            crop.unsqueeze(0),  # Add batch dim
            size=(self.resize_to, self.resize_to),
            mode='bilinear',
            align_corners=False
        ).squeeze(0)  # [3, resize_to, resize_to]

        return {
            "pixel_values": crop,
            "input_ids": batch["input_ids"],
            "attention_mask": batch["attention_mask"],
            "imgid": batch["imgid"],
        }
        

class DebatchDataset(Dataset):
    """no-caption dataset: processes single images on-the-fly with empty text prompt."""
    
    def __init__(self, dataset, processor):
        self.dataset = dataset
        if "blip" in processor:
            print("Using Blip processor")
            self.processor = BlipProcessor.from_pretrained(processor)
            self.prompt = "Describe this image"
        elif "llava" in processor:
            print("Using Llava processor")
            self.processor = LlavaProcessor.from_pretrained(processor)
            self.prompt = "USER: <image>\nDescribe this image. \nASSISTANT:"
        else:
            raise ValueError(f"Unsupported processor: {processor}")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        raw_item = self.dataset[idx]
        image = raw_item["image"]

        # Process exactly like batched version but with B=1 and prompt="Describe this image"
        processed = self.processor(
            images=image,
            text=self.prompt, 
            return_tensors="pt",
            padding=True,
        )

        return {
            "pixel_values": processed["pixel_values"],
            "input_ids": processed["input_ids"],
            "attention_mask": processed["attention_mask"],
            "imgid": raw_item["imgid"],
        }