import json
import os
from typing import Dict, List, Optional, Tuple

import torch

from .model import BitTransformerLM
from .training import train_loop


def collapse_submodel(
    cluster_data: List[List[int]],
    target_params: Dict,
    floors: Optional[Dict[str, float]] = None,
    max_rounds: int = 3,
    width_scale: float = 1.5,
    forward_kwargs: Optional[Dict] = None,
) -> Tuple[BitTransformerLM, Dict[str, float]]:
    """Distill a submodel from clustered bit sequences.

    The routine deepens the target model when telemetry floors are unmet and,
    after the first deepening fails, widens the hidden dimensions by
    ``width_scale`` once before retrying. Returns the distilled model and its
    final telemetry metrics.
    """
    if floors is None:
        floors = {"negentropy": 0.5, "lz_complexity": 0.3, "symbiosis_score": 0.5}

    bit_tensor = torch.tensor(cluster_data, dtype=torch.long)
    n = len(bit_tensor)
    split = max(1, int(0.8 * n))
    train_bits = bit_tensor[:split]
    val_bits = bit_tensor[split:]
    if len(val_bits) == 0:
        val_bits = train_bits

    params = target_params.copy()
    metrics: Dict[str, float] = {}
    width_scaled = False
    for round_idx in range(max_rounds):
        model = BitTransformerLM(**params)
        train_loop(
            model,
            train_bits,
            epochs=2,
            compress_prob=0.5,
            direct_prob=0.0,
            log=False,
            forward_kwargs=forward_kwargs,
        )
        with torch.no_grad():
            logits, telemetry = model(val_bits, **(forward_kwargs or {}))
            neg_k = model.negentropy_logits(logits).mean().item()
            lz_c = model.lz_complexity_logits(logits).mean().item()
            sym_s = telemetry["symbiosis_score"].mean().item()
        metrics = {
            "negentropy": neg_k,
            "lz_complexity": lz_c,
            "symbiosis_score": sym_s,
        }
        if (
            neg_k >= floors["negentropy"]
            and lz_c >= floors["lz_complexity"]
            and sym_s >= floors["symbiosis_score"]
        ):
            break
        if round_idx == 0:
            params["num_layers"] = max(1, params.get("num_layers", 1)) + 1
        elif not width_scaled:
            params["d_model"] = int(params.get("d_model", 32) * width_scale)
            params["dim_feedforward"] = int(
                params.get("dim_feedforward", 64) * width_scale
            )
            width_scaled = True
        else:
            params["num_layers"] = max(1, params.get("num_layers", 1)) + 1
    return model, metrics


def save_distilled_model(
    model: BitTransformerLM,
    path: str,
    metrics: Dict[str, float],
    floors: Optional[Dict[str, float]] = None,
) -> None:
    """Serialize a distilled model and its metric summary to disk.

    Weights are written to ``path`` and a ``metrics.json`` file is placed in the
    same directory containing the achieved metrics alongside the target floors.
    """
    torch.save(model.state_dict(), path)
    payload = {"metrics": metrics, "floors": floors or {}}
    metrics_path = os.path.join(os.path.dirname(path), "metrics.json")
    with open(metrics_path, "w") as f:
        json.dump(payload, f)