"""
Instruction Residuals for LLM Continuous Pre-Training

Based on "Balancing Continuous Pre-Training and Instruction Fine-Tuning" 
(Samsung Research, 2024) and "Editing Models with Task Arithmetic" (Ilharco et al., 2022)

Key findings from the literature:
- Instruction capabilities are portable across models from the same family
- Simple element-wise addition (task arithmetic) is effective for merging
- No additional scaling or normalization is typically needed
- Tokenizer alignment is critical before applying residuals
"""

import hashlib
import json
import os
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from typing import Dict, Optional, Tuple, Union

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def _align_tokenizer_and_embeddings(
    base_model: AutoModelForCausalLM, base_tokenizer: AutoTokenizer
) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
    """
    Align tokenizer and embeddings by adding PAD token if missing.

    Critical step from apply_delta.py scripts (StableVicuna, LLaVA):
    - Ensures tokenizer compatibility
    - Resizes embeddings to match vocabulary
    - Zeros newly added embedding rows to prevent contamination

    Args:
        base_model: Model to align
        base_tokenizer: Tokenizer to align

    Returns:
        Tuple of (aligned_model, aligned_tokenizer)
    """
    DEFAULT_PAD_TOKEN = "[PAD]"
    num_new = 0

    if base_tokenizer.pad_token is None:
        num_new = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))

    if num_new > 0:
        base_model.resize_token_embeddings(len(base_tokenizer))

        # Zero out new embedding rows to avoid random initialization
        input_embeddings = base_model.get_input_embeddings().weight.data
        output_embeddings = base_model.get_output_embeddings().weight.data

        input_embeddings[-num_new:] = 0
        output_embeddings[-num_new:] = 0

    return base_model, base_tokenizer


@dataclass
class ResidualsConfig:
    residuals_version: str
    format: str
    dtype: str
    param_count: int
    shapes_hash: str
    parameter_names_hash: str
    tokenizer_name: Optional[str]
    created_at: str
    library: Dict[str, str]


class Residuals:
    """Container for instruction residuals with load/save/apply ergonomics."""

    def __init__(self, state_dict: Dict[str, torch.Tensor], config: ResidualsConfig, instruct_tokenizer: Optional[AutoTokenizer] = None):
        self.state_dict = state_dict
        self.config = config
        self.instruct_tokenizer = instruct_tokenizer

    @staticmethod
    def from_models(
        base_model: Optional[AutoModelForCausalLM] = None,
        instruct_model: Optional[AutoModelForCausalLM] = None,
        base_model_name: Optional[str] = None,
        instruct_model_name: Optional[str] = None,
        instruct_tokenizer: Optional[AutoTokenizer] = None,
        instruct_tokenizer_name: Optional[str] = None,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = "cpu",
    ) -> "Residuals":
        # Allow passing either models or names; load if names were provided
        if base_model is None:
            if base_model_name is None:
                raise ValueError("Either base_model or base_model_name must be provided")
            base_model = AutoModelForCausalLM.from_pretrained(
                base_model_name, torch_dtype=dtype
            )
            # Place on requested device if provided
            try:
                if device is not None:
                    base_model = base_model.to(device)
                base_model = base_model.to(dtype=dtype)
            except Exception:
                pass
        if instruct_model is None:
            if instruct_model_name is None:
                raise ValueError("Either instruct_model or instruct_model_name must be provided")
            instruct_model = AutoModelForCausalLM.from_pretrained(
                instruct_model_name, torch_dtype=dtype
            )
            try:
                if device is not None:
                    instruct_model = instruct_model.to(device)
                instruct_model = instruct_model.to(dtype=dtype)
            except Exception:
                pass

        base_sd = base_model.state_dict()
        inst_sd = instruct_model.state_dict()
        state: Dict[str, torch.Tensor] = {}
        for key, base_param in base_sd.items():
            if key not in inst_sd:
                raise KeyError(f"Missing key in instruct model: {key}")
            inst_param = inst_sd[key]
            if base_param.shape != inst_param.shape:
                raise ValueError(
                    f"Shape mismatch for {key}: base {base_param.shape} vs instruct {inst_param.shape}"
                )
            state[key] = (inst_param - base_param).to(dtype=inst_param.dtype)
        # Determine instruct tokenizer
        tok = instruct_tokenizer
        if tok is None and instruct_tokenizer_name is not None:
            tok = AutoTokenizer.from_pretrained(instruct_tokenizer_name, use_fast=False)

        tok_name = instruct_tokenizer_name
        cfg = _build_config(state, tokenizer_name=tok_name)
        return Residuals(state, cfg, instruct_tokenizer=tok)

    @staticmethod
    def from_pretrained(
        path: str,
        map_location: Union[str, torch.device] = "cpu",
    ) -> "Residuals":
        residuals_path = os.path.join(path, "pytorch_model.bin")
        state: Dict[str, torch.Tensor] = torch.load(residuals_path, map_location=map_location)
        cfg_path = os.path.join(path, "config.json")
        if os.path.exists(cfg_path):
            with open(cfg_path, "r", encoding="utf-8") as f:
                raw = json.load(f)
            cfg = ResidualsConfig(**raw)
        else:
            # No fallback; tokenizer must have been saved in this format
            raise FileNotFoundError("config.json not found alongside residuals; tokenizer required by current format")
        # Load tokenizer saved in the same directory
        tok = AutoTokenizer.from_pretrained(path, use_fast=False)
        return Residuals(state, cfg, instruct_tokenizer=tok)

    def apply(
        self,
        base_model: AutoModelForCausalLM,
        base_tokenizer: Optional[AutoTokenizer] = None,
        instruct_tokenizer: Optional[AutoTokenizer] = None,
        out_dir: Optional[str] = None,
        scale: float = 1.0,
    ) -> AutoModelForCausalLM:
        # Prefer provided instruct_tokenizer; else fall back to stored one
        if instruct_tokenizer is None and self.instruct_tokenizer is not None:
            instruct_tokenizer = self.instruct_tokenizer

        if base_tokenizer is not None and instruct_tokenizer is not None:
            base_model, base_tokenizer = _align_tokenizer_and_embeddings(
                base_model, base_tokenizer
            )

        groups: Dict[int, Dict[str, torch.Tensor]] = {}
        tensors_by_ptr: Dict[int, torch.Tensor] = {}

        for n, p in base_model.named_parameters():
            if n not in self.state_dict:
                raise KeyError(f"Residuals missing key: {n}")
            delta = self.state_dict[n]
            if delta.shape != p.shape:
                raise ValueError(
                    f"Shape mismatch for {n}: param {p.shape} vs delta {delta.shape}"
                )
            try:
                ptr = p.data_ptr()
            except Exception:
                ptr = id(p)
            if ptr not in groups:
                groups[ptr] = {}
                tensors_by_ptr[ptr] = p
            groups[ptr][n] = delta

        with torch.no_grad():
            for ptr, name_to_delta in groups.items():
                t = tensors_by_ptr[ptr]
                total_delta = None
                for d in name_to_delta.values():
                    d_ = (d * scale).to(t.dtype)
                    total_delta = d_ if total_delta is None else (total_delta + d_)
                t.add_(total_delta)

        for n, b in base_model.named_buffers():
            if n in self.state_dict:
                d = self.state_dict[n]
                if d.shape != b.shape:
                    raise ValueError(
                        f"Shape mismatch for buffer {n}: param {b.shape} vs delta {d.shape}"
                    )
                with torch.no_grad():
                    b.add_((d * scale).to(b.dtype))

        if out_dir:
            base_model.save_pretrained(out_dir)
            if instruct_tokenizer is not None:
                instruct_tokenizer.save_pretrained(out_dir)

        return base_model

    def save_pretrained(self, out_dir: str) -> None:
        os.makedirs(out_dir, exist_ok=True)
        torch.save(self.state_dict, os.path.join(out_dir, "pytorch_model.bin"))
        # Always require tokenizer to be saved
        if self.instruct_tokenizer is None:
            raise ValueError("instruct_tokenizer must be provided to save_pretrained; set it via from_models(... instruct_tokenizer=... or instruct_tokenizer_name=...)")
        self.instruct_tokenizer.save_pretrained(out_dir)
        with open(os.path.join(out_dir, "config.json"), "w", encoding="utf-8") as f:
            json.dump(asdict(self.config), f, indent=2)

    def to(
        self,
        device: Optional[Union[str, torch.device]] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> "Residuals":
        """Return a new Residuals instance with tensors moved/cast to device/dtype.

        Does not modify the original instance.
        """
        new_state: Dict[str, torch.Tensor] = {}
        for k, v in self.state_dict.items():
            new_state[k] = v.to(device=device if device is not None else v.device, dtype=dtype if dtype is not None else v.dtype)

        # Update config dtype string based on new tensors
        names = sorted(new_state.keys())
        dtypes = [str(new_state[n].dtype).split(".")[-1] for n in names]
        dtype_str = max(set(dtypes), key=dtypes.count) if dtypes else "float32"
        new_cfg = ResidualsConfig(
            residuals_version=self.config.residuals_version,
            format=self.config.format,
            dtype=dtype_str,
            param_count=self.config.param_count,
            shapes_hash=self.config.shapes_hash,
            parameter_names_hash=self.config.parameter_names_hash,
            tokenizer_name=self.config.tokenizer_name,
            created_at=self.config.created_at,
            library=self.config.library,
        )
        return Residuals(new_state, new_cfg, instruct_tokenizer=self.instruct_tokenizer)


def _build_config(
    residuals: Dict[str, torch.Tensor], tokenizer_name: Optional[str]
) -> ResidualsConfig:
    names = sorted(residuals.keys())
    shapes = [tuple(residuals[n].shape) for n in names]
    names_hash = hashlib.sha256("|".join(names).encode()).hexdigest()
    shapes_hash = hashlib.sha256(
        "|".join([f"{n}:{list(s)}" for n, s in zip(names, shapes)]).encode()
    ).hexdigest()
    # Dtype may vary across tensors; record the most common dtype or 'mixed'
    dtypes = [str(t.dtype).split(".")[-1] for t in residuals.values()]
    dtype = max(set(dtypes), key=dtypes.count) if dtypes else "float32"
    return ResidualsConfig(
        residuals_version="1.0",
        format="pytorch",
        dtype=dtype,
        param_count=len(residuals),
        shapes_hash=shapes_hash,
        parameter_names_hash=names_hash,
        tokenizer_name=tokenizer_name,
        created_at=datetime.now(timezone.utc).isoformat(),
        library={"package": "residuals"},
    )
