"""
Derived from Andrej Karpathy's nanochat project.

MIT License

Copyright (c) 2025 Andrej Karpathy

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"""

from __future__ import annotations

import argparse

from dropout_decay.specs import DropoutCondition, ModelSpec


DEFAULT_DROPOUT_RATES = [0.0, 0.02, 0.05, 0.08, 0.10, 0.14, 0.20, 0.30, 0.50]


def clean_name(value: str) -> str:
    cleaned = "".join(
        c if c.isalnum() or c in {"-", "_"} else "_" for c in value.strip()
    )
    return cleaned.strip("_") or "model"


def rate_label(rate: float) -> str:
    label = f"{rate:.3f}".rstrip("0").rstrip(".")
    return label if label else "0"


def parse_model_spec(raw: str) -> ModelSpec:
    if "=" in raw:
        name, dims = raw.split("=", 1)
        name = clean_name(name)
    else:
        dims = raw
        name = ""
    parts = dims.lower().replace(",", "x").split("x")
    if len(parts) != 3:
        raise argparse.ArgumentTypeError(
            "model specs must look like 8x8x256 or name=8x8x256"
        )
    try:
        n_layer, n_head, n_embd = [int(part) for part in parts]
    except ValueError as exc:
        raise argparse.ArgumentTypeError("model dimensions must be integers") from exc
    if n_layer <= 0 or n_head <= 0 or n_embd <= 0:
        raise argparse.ArgumentTypeError("model dimensions must be positive")
    if n_embd % n_head != 0:
        raise argparse.ArgumentTypeError("n_embd must be divisible by n_head")
    if (n_embd // n_head) % 2 != 0:
        raise argparse.ArgumentTypeError("n_embd / n_head must be even for rotary attention")
    return ModelSpec(
        name or f"L{n_layer}_H{n_head}_D{n_embd}",
        n_layer,
        n_head,
        n_embd,
    )


def parse_decay_spec(raw: str) -> DropoutCondition:
    parts = raw.split(":")
    if len(parts) not in {3, 4, 5}:
        raise argparse.ArgumentTypeError(
            "decay specs must look like "
            "name:initial:final[:cosine|smoothstep|linear[:decay_tokens]]"
        )
    name = clean_name(parts[0])
    try:
        initial = float(parts[1])
        final = float(parts[2])
        decay_tokens = int(parts[4]) if len(parts) == 5 else None
    except ValueError as exc:
        raise argparse.ArgumentTypeError(
            "decay dropout values and decay_tokens must be numeric"
        ) from exc
    schedule = parts[3] if len(parts) >= 4 else "cosine"
    if schedule not in {"cosine", "smoothstep", "linear"}:
        raise argparse.ArgumentTypeError(
            "decay schedule must be cosine, smoothstep, or linear"
        )
    return DropoutCondition(
        name=name,
        kind="decay",
        initial=initial,
        final=final,
        schedule=schedule,
        decay_tokens=decay_tokens,
    )


def parse_anchor_decay_spec(raw: str) -> DropoutCondition:
    if ":" not in raw:
        raise argparse.ArgumentTypeError(
            "anchor decay specs must look like name:250000=0.60,500000=0.40"
        )
    name, raw_anchors = raw.split(":", 1)
    anchors: list[tuple[int, float]] = []
    for piece in raw_anchors.split(","):
        if "=" not in piece:
            raise argparse.ArgumentTypeError(
                "anchor decay anchors must look like token_count=dropout"
            )
        raw_tokens, raw_dropout = piece.split("=", 1)
        try:
            token_count = int(raw_tokens)
            dropout = float(raw_dropout)
        except ValueError as exc:
            raise argparse.ArgumentTypeError(
                "anchor token counts must be integers and dropout values numeric"
            ) from exc
        if token_count <= 0:
            raise argparse.ArgumentTypeError("anchor token counts must be positive")
        if not 0.0 <= dropout < 1.0:
            raise argparse.ArgumentTypeError("anchor dropout must satisfy 0 <= p < 1")
        anchors.append((token_count, dropout))
    anchors = sorted(set(anchors))
    if len(anchors) < 2:
        raise argparse.ArgumentTypeError("provide at least two dropout anchors")
    for left, right in zip(anchors, anchors[1:]):
        if right[1] > left[1]:
            raise argparse.ArgumentTypeError(
                "anchor dropout values must be non-increasing as token counts grow"
            )
    return DropoutCondition(
        name=clean_name(name),
        kind="anchor_decay",
        initial=anchors[0][1],
        final=anchors[-1][1],
        schedule="log_prefix_anchor",
        anchors=tuple(anchors),
    )


def default_seeds(mode: str, seeds: list[int] | None) -> list[int]:
    if seeds:
        return seeds
    return [1] if mode == "screen_static" else [1, 2, 3]