from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path

import hashlib
import numpy as np
import torch


"""
Core types for BCT
"""


@dataclass
class BlockType:
    """Specifies a type of block like text or audio. This should be serializable."""

    name: str
    is_causal: bool
    chunk_size: Optional[int] = None
    description: Optional[str] = None

    def __str__(self):
        name_hash = hashlib.md5(self.name.encode()).hexdigest()
        color_code = f"#{name_hash[:6]}"  # Use first 6 chars of hash for hex color
        colored_name = f"\033[38;2;{int(color_code[1:3], 16)};{int(color_code[3:5], 16)};{int(color_code[5:7], 16)}m{self.name}\033[0m"
        return f"{colored_name} ({'causal' if self.is_causal else 'non-causal'}) {f'chunk={self.chunk_size}' if self.chunk_size else ''}"

    def to_dict(self) -> dict:
        return {"name": self.name, "is_causal": self.is_causal}

    @classmethod
    def from_dict(cls, data: dict) -> "BlockType":
        return cls(name=data["name"], is_causal=data["is_causal"])


@dataclass
class Block:
    """An instance of a block with data"""

    spec: BlockType
    inputs: dict[str, torch.Tensor] = field(default_factory=dict)  # (T, D)
    targets: dict[str, torch.Tensor] = field(default_factory=dict)  # (T, D)
    debug_text: Optional[str] = None  # Optional debug text for printing

    def __len__(self):
        """Length of the block in tokens"""
        if len(self.inputs) == 0:
            return 0
        return next(iter(self.inputs.values())).shape[0]

    @staticmethod
    def shift_left(
        x: torch.Tensor | np.ndarray, pad_token: int, n: int = 1
    ) -> torch.Tensor | np.ndarray:
        """Shift the tensor to the left by n steps. Assume first dimension is time."""
        x = x.clone()
        x[:-n] = x.clone()[n:]
        x[-n:] = pad_token
        return x

    def __str__(self):
        """Multi-line string representation of the block"""

        # Helper function to format tensor dict showing only first and last 3 elements
        def format_tensor_dict(tensor_dict):
            if not tensor_dict:
                return ""

            result = ""
            for key, tensor in tensor_dict.items():
                dims = tensor.shape[0]
                dims_str = f"({dims} dims)" if dims > 1 else ""
                if tensor.ndim == 2:
                    if tensor.shape[0] > 1:
                        tensor = tensor.mean(dim=0)
                    else:
                        tensor = tensor[0]

                if tensor.shape[0] <= 6:
                    # If tensor has 6 or fewer elements, show all
                    tensor_str = str(tensor.tolist())
                    result += f"\n    {key}: {dims_str} \033[2m{tensor_str}\033[0m"
                else:
                    # Show first 3 and last 3 elements
                    first_three = tensor[:3].tolist()
                    first_three_str = ", ".join(f"{x}" for x in first_three)
                    last_three = tensor[-3:].tolist()
                    last_three_str = ", ".join(f"{x}" for x in last_three)
                    result += f"\n    {key}: {dims_str} \033[2m[{first_three_str}, ..., {last_three_str}]\033[0m"

            return result

        inputs_str = format_tensor_dict(self.inputs)
        targets_str = format_tensor_dict(self.targets)
        s = f"{self.spec} ({len(self)})"
        if self.debug_text:
            s += f" \033[96m'{self.debug_text}'\033[0m"  # Cyan color for debug text
        if inputs_str:
            s += f"{inputs_str}"
        if targets_str:
            s += f"{targets_str}"
        return s


@dataclass
class BlockSequence:
    """A sequence of blocks representing a single sample"""

    blocks: list[Block]

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

    def __getitem__(self, idx: int) -> Block:
        return self.blocks[idx]

    def __setitem__(self, idx: int, value: Block):
        self.blocks[idx] = value

    def __iter__(self):
        return iter(self.blocks)

    def __add__(self, other: "BlockSequence") -> "BlockSequence":
        return BlockSequence(self.blocks + other.blocks)

    def __str__(self):
        """Multi-line string representation of the block sequence"""
        s = "BlockSequence(\n"
        for block in self.blocks:
            s += f"  {block}\n"
        s += ")"
        return s

    @property
    def n_tokens(self) -> int:
        return sum(len(block) for block in self.blocks)

    def crop_to_max_tokens(self, max_tokens: int) -> "BlockSequence":
        """Crop the blocks to the maximum number of tokens."""
        total_tokens = self.n_tokens
        assert max_tokens <= total_tokens, (max_tokens, total_tokens)

        result_blocks = []
        for block in self.blocks:
            # Calculate the current length of all blocks in the result
            cur_len = sum(len(b) for b in result_blocks)

            # If adding this block would exceed max_tokens, we need to truncate it
            if cur_len + len(block) >= max_tokens:
                # Truncate each input and target tensor to fit within max_tokens
                block.inputs = {k: v[: max_tokens - cur_len] for k, v in block.inputs.items()}
                block.targets = {k: v[: max_tokens - cur_len] for k, v in block.targets.items()}
                result_blocks.append(block)
                break
            result_blocks.append(block)

        result = BlockSequence(result_blocks)
        assert result.n_tokens == max_tokens, (result.n_tokens, max_tokens)
        return result

    def save(self, path: Path | str) -> None:
        """Save BlockSequence to a single PyTorch file."""
        path = Path(path)

        # Build tensors dict with unique keys
        tensors = {}
        metadata = {"version": "1.0", "type": "BlockSequence", "blocks": []}

        for block_idx, block in enumerate(self.blocks):
            block_meta = {"spec": block.spec.to_dict(), "inputs": {}, "targets": {}}

            # Save input tensors
            for key, tensor in block.inputs.items():
                tensor_key = f"block_{block_idx}_inputs_{key}"
                tensors[tensor_key] = tensor
                block_meta["inputs"][key] = tensor_key

            # Save target tensors
            for key, tensor in block.targets.items():
                tensor_key = f"block_{block_idx}_targets_{key}"
                tensors[tensor_key] = tensor
                block_meta["targets"][key] = tensor_key

            metadata["blocks"].append(block_meta)

        # Save everything to single file
        torch.save({"metadata": metadata, "tensors": tensors}, path)

    @classmethod
    def load(cls, path: Path | str) -> "BlockSequence":
        """Load BlockSequence from a PyTorch file."""
        path = Path(path)
        data = torch.load(path, map_location="cpu", weights_only=False)

        metadata = data["metadata"]
        tensors = data["tensors"]

        blocks = []
        for block_meta in metadata["blocks"]:
            # Reconstruct BlockType
            spec = BlockType.from_dict(block_meta["spec"])

            # Reconstruct input tensors
            inputs = {}
            for key, tensor_key in block_meta["inputs"].items():
                inputs[key] = tensors[tensor_key]

            # Reconstruct target tensors
            targets = {}
            for key, tensor_key in block_meta["targets"].items():
                targets[key] = tensors[tensor_key]

            blocks.append(Block(spec=spec, inputs=inputs, targets=targets))

        return cls(blocks=blocks)


@dataclass
class PackedBlockSequence:
    """Multiple block sequences packed together, representing a single batch"""

    block_sequences: list[BlockSequence]
    avg_seq_length_before_crop: float = 0.0

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

    def __getitem__(self, idx: int) -> BlockSequence:
        return self.block_sequences[idx]

    def __setitem__(self, idx: int, value: BlockSequence):
        self.block_sequences[idx] = value

    def __iter__(self):
        return iter(self.block_sequences)

    def __add__(self, other: "PackedBlockSequence") -> "PackedBlockSequence":
        return PackedBlockSequence(self.block_sequences + other.block_sequences)

    def append(self, block_sequence: BlockSequence):
        self.block_sequences.append(block_sequence)

    @property
    def n_tokens(self) -> int:
        return sum(block_seq.n_tokens for block_seq in self.block_sequences)

    def crop_to_max_tokens(self, max_tokens: int) -> "PackedBlockSequence":
        """Crop the blocks to the maximum number of tokens."""
        result_block_sequences = PackedBlockSequence([])
        for block_sequence in self.block_sequences:
            cur_len = result_block_sequences.n_tokens
            if cur_len + block_sequence.n_tokens >= max_tokens:
                result_block_sequences.append(block_sequence.crop_to_max_tokens(max_tokens - cur_len))
                break
            result_block_sequences.append(block_sequence)

        assert result_block_sequences.n_tokens == max_tokens
        result_block_sequences.avg_seq_length_before_crop = self.avg_seq_length_before_crop
        return result_block_sequences

    def save(self, path: Path | str) -> None:
        """Save PackedBlockSequence to a single PyTorch file."""
        path = Path(path)

        # Build tensors dict with unique keys
        tensors = {}
        metadata = {"version": "1.0", "type": "PackedBlockSequence", "block_sequences": []}

        for seq_idx, block_sequence in enumerate(self.block_sequences):
            seq_meta = {"blocks": []}

            for block_idx, block in enumerate(block_sequence.blocks):
                block_meta = {"spec": block.spec.to_dict(), "inputs": {}, "targets": {}}

                # Save input tensors
                for key, tensor in block.inputs.items():
                    tensor_key = f"seq_{seq_idx}_block_{block_idx}_inputs_{key}"
                    tensors[tensor_key] = tensor
                    block_meta["inputs"][key] = tensor_key

                # Save target tensors
                for key, tensor in block.targets.items():
                    tensor_key = f"seq_{seq_idx}_block_{block_idx}_targets_{key}"
                    tensors[tensor_key] = tensor
                    block_meta["targets"][key] = tensor_key

                seq_meta["blocks"].append(block_meta)

            metadata["block_sequences"].append(seq_meta)

        # Save everything to single file
        torch.save({"metadata": metadata, "tensors": tensors}, path)

    @classmethod
    def load(cls, path: Path | str) -> "PackedBlockSequence":
        """Load PackedBlockSequence from a PyTorch file."""
        path = Path(path)
        data = torch.load(path, map_location="cpu", weights_only=False)

        metadata = data["metadata"]
        tensors = data["tensors"]

        block_sequences = []
        for seq_meta in metadata["block_sequences"]:
            blocks = []
            for block_meta in seq_meta["blocks"]:
                # Reconstruct BlockType
                spec = BlockType.from_dict(block_meta["spec"])

                # Reconstruct input tensors
                inputs = {}
                for key, tensor_key in block_meta["inputs"].items():
                    inputs[key] = tensors[tensor_key]

                # Reconstruct target tensors
                targets = {}
                for key, tensor_key in block_meta["targets"].items():
                    targets[key] = tensors[tensor_key]

                blocks.append(Block(spec=spec, inputs=inputs, targets=targets))

            block_sequences.append(BlockSequence(blocks=blocks))

        return cls(block_sequences=block_sequences)


from collections import Counter, deque


class BlockUsageStatistics:
    """Statistics about training data using a rolling window"""

    def __init__(self, window_size: int = 100):
        self.window_size = window_size
        self.window = deque(maxlen=window_size)  # Stores (packed_block_sequence, per_block_losses)

        # Statistics (computed lazily when accessed)
        self.block_counter = Counter()
        self.block_pair_counter = Counter()
        self.block_token_counter = Counter()
        self.block_loss_sum = Counter()
        self.block_loss_count = Counter()
        self.total_samples = 0
        self._dirty = False  # Flag to track if recompute is needed

    def update(
        self, packed_block_sequence: PackedBlockSequence, per_block_losses: Optional[list] = None
    ):
        # Add new sample to window (automatically removes oldest if window is full)
        self.window.append((packed_block_sequence, per_block_losses))
        self._dirty = True  # Mark statistics as stale

    def _ensure_fresh(self):
        """Recompute statistics if they are stale"""
        if not self._dirty:
            return

        # Reset counters
        self.block_counter = Counter()
        self.block_pair_counter = Counter()
        self.total_samples = 0
        self.block_token_counter = Counter()
        self.block_loss_sum = Counter()
        self.block_loss_count = Counter()

        # Recompute from window
        for packed_block_sequence, per_block_losses in self.window:
            self.total_samples += len(packed_block_sequence)

            # If per_block_losses is provided, use the actual per-block loss values
            if per_block_losses is not None:
                # per_block_losses is a list of tuples: (sample, block_type, loss_value)
                for sample, _, loss_value in per_block_losses:
                    conditioning_blocks = sample[:-1]  # all conditioning blocks

                    # Check for special cases
                    if len(conditioning_blocks) == 0:
                        # No conditioning blocks - track as "no_conditioning"
                        self.block_loss_sum["no_conditioning"] += loss_value
                        self.block_loss_count["no_conditioning"] += 1
                    elif len(conditioning_blocks) == 1:
                        # Only one conditioning block - track as "only_{block_name}"
                        block_name = conditioning_blocks[0].spec.name
                        only_key = f"only_{block_name}"
                        self.block_loss_sum[only_key] += loss_value
                        self.block_loss_count[only_key] += 1

                    # Track individual conditioning blocks
                    for block in conditioning_blocks:
                        self.block_loss_sum[block.spec.name] += loss_value
                        self.block_loss_count[block.spec.name] += 1

            for block_sequence in packed_block_sequence:
                for block in block_sequence:
                    self.block_counter[block.spec.name] += 1
                    self.block_token_counter[block.spec.name] += len(block)

                    for other_block in block_sequence:
                        if block.spec.name == other_block.spec.name:
                            continue
                        pair = tuple(sorted((block.spec.name, other_block.spec.name)))
                        self.block_pair_counter[pair] += 1

        self._dirty = False

    def __str__(self):
        """Summary of the block usage statistics"""
        self._ensure_fresh()

        # normalize the counters
        block_usage_stats = {k: v / self.total_samples for k, v in self.block_counter.items()}

        total_tokens = sum(self.block_token_counter.values())
        block_token_usage_stats = {k: v / total_tokens for k, v in self.block_token_counter.items()}

        block_usage_stats = dict(sorted(block_usage_stats.items(), key=lambda x: x[1], reverse=True))
        block_token_usage_stats = dict(
            sorted(block_token_usage_stats.items(), key=lambda x: x[1], reverse=True)
        )

        # Format the stats as a readable string
        blocks_str = " | ".join(f"{k}={v:.3f}" for k, v in block_usage_stats.items())
        tokens_str = " | ".join(f"{k}={v:.3f}" for k, v in block_token_usage_stats.items())

        result = f"Blocks/sample: {blocks_str}\n% tokens: {tokens_str}"

        # Add loss statistics if available
        if self.block_loss_count:
            block_loss_stats = {
                k: self.block_loss_sum[k] / self.block_loss_count[k]
                for k in self.block_loss_count.keys()
            }
            block_loss_stats = dict(sorted(block_loss_stats.items(), key=lambda x: x[1], reverse=True))
            loss_str = " | ".join(f"{k}={v:.4f}" for k, v in block_loss_stats.items())
            result += f"\nAvg loss: {loss_str}"

        return result

    def get_wandb_metrics(self) -> dict:
        """Get block loss metrics formatted for wandb logging"""
        self._ensure_fresh()
        metrics = {}

        if self.block_loss_count:
            # Add average loss per block type
            for block_type in self.block_loss_count.keys():
                avg_loss = self.block_loss_sum[block_type] / self.block_loss_count[block_type]
                metrics[f"block_loss/{block_type}"] = avg_loss

        # Add usage statistics
        if self.total_samples > 0:
            for block_type, count in self.block_counter.items():
                blocks_per_sample = count / self.total_samples
                metrics[f"block_usage/blocks_per_sample/{block_type}"] = blocks_per_sample

            total_tokens = sum(self.block_token_counter.values())
            if total_tokens > 0:
                for block_type, tokens in self.block_token_counter.items():
                    token_fraction = tokens / total_tokens
                    metrics[f"block_usage/token_fraction/{block_type}"] = token_fraction

        return metrics