#!/usr/bin/env python3
"""Build APE weights on the frozen Florence-Bridge valid-anchor support."""

from __future__ import annotations

import argparse
import json
from collections import OrderedDict, defaultdict
from pathlib import Path

import numpy as np
import pyarrow.parquet as parquet
import torch
from torch.utils.data import DataLoader, TensorDataset

from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest

from external_baselines.keyframe_focused_adapted.keyframe_cache import (
    ActionCopycat,
    history_target,
    q99_normalize,
)


def episode_path(data_root: Path, dataset_name: str, trajectory_id: int) -> Path:
    return (
        data_root
        / dataset_name
        / "data"
        / f"chunk-{trajectory_id // 1000:03d}"
        / f"episode_{trajectory_id:06d}.parquet"
    )


def read_actions(
    data_root: Path,
    dataset_name: str,
    trajectory_id: int,
    action_dim: int | None,
) -> np.ndarray:
    path = episode_path(data_root, dataset_name, trajectory_id)
    values = parquet.read_table(path, columns=["action"])["action"].to_pylist()
    actions = np.asarray(values, dtype=np.float32)
    if actions.ndim != 2 or (action_dim is not None and actions.shape[1] != action_dim):
        raise RuntimeError(f"unexpected action shape {actions.shape} in {path}")
    return actions


def resolve_statistics_key(
    dataset_name: str,
    statistics_payload: dict,
    explicit_mapping: dict[str, str],
) -> str:
    """Resolve Bridge aliases and arbitrary RoboTwin dataset statistics."""

    aliases = {
        "bridge_orig_lerobot": "oxe_bridge",
        "fractal20220817_data_lerobot": "oxe_rt1",
    }
    candidates = (
        explicit_mapping.get(dataset_name),
        aliases.get(dataset_name),
        dataset_name,
    )
    for candidate in candidates:
        if candidate and candidate in statistics_payload:
            return candidate
    if len(statistics_payload) == 1:
        return next(iter(statistics_payload))
    raise KeyError(
        f"no statistics entry for {dataset_name!r}; provide --statistics-key-map"
    )


def requested_support(manifest, args):
    if args.scope == "full":
        return {
            dataset.name: OrderedDict(
                (int(episode.trajectory_id), tuple(int(x) for x in episode.anchors))
                for episode in dataset.episodes
            )
            for dataset in manifest.datasets
        }
    selected = defaultdict(lambda: defaultdict(set))
    for draw_id in range(args.draw_start, args.draw_start + args.draw_count):
        trace = manifest.sample(args.draw_seed, LogicalDraw(0, draw_id))
        selected[trace.dataset_name][int(trace.trajectory_id)].add(int(trace.base_index))
    return {
        dataset.name: OrderedDict(
            (
                int(episode.trajectory_id),
                tuple(sorted(selected[dataset.name][int(episode.trajectory_id)])),
            )
            for episode in dataset.episodes
            if selected[dataset.name].get(int(episode.trajectory_id))
        )
        for dataset in manifest.datasets
    }


def fitting_support(manifest, max_anchors: int | None):
    result = {}
    for dataset in manifest.datasets:
        episodes = OrderedDict()
        count = 0
        for episode in dataset.episodes:
            anchors = tuple(int(x) for x in episode.anchors if int(x) >= 9)
            if max_anchors is not None:
                anchors = anchors[: max(0, max_anchors - count)]
            if anchors:
                episodes[int(episode.trajectory_id)] = anchors
                count += len(anchors)
            if max_anchors is not None and count >= max_anchors:
                break
        if count == 0:
            raise RuntimeError(f"no copycat fitting anchors for {dataset.name}")
        result[dataset.name] = episodes
    return result


def examples_for_support(
    data_root,
    dataset_name,
    support,
    statistics,
    *,
    action_dim,
    binary_action_indices,
    history=9,
):
    inputs, targets = [], []
    for trajectory_id, anchors in support.items():
        actions = q99_normalize(
            read_actions(data_root, dataset_name, trajectory_id, action_dim),
            statistics,
            binary_action_indices=binary_action_indices,
        )
        for anchor in anchors:
            x, y = history_target(actions, anchor, history=history)
            inputs.append(x)
            targets.append(y)
    return np.asarray(inputs, dtype=np.float32), np.asarray(targets, dtype=np.float32)


def train_copycat(inputs, targets, *, action_dim, epochs, device, seed=0):
    generator = torch.Generator().manual_seed(seed)
    permutation = torch.randperm(len(inputs), generator=generator)
    split = int(len(inputs) * 0.8)
    train_indices = permutation[:split]
    val_indices = permutation[split:]
    x = torch.from_numpy(inputs)
    y = torch.from_numpy(targets)
    train = TensorDataset(x[train_indices], y[train_indices])
    val = TensorDataset(x[val_indices], y[val_indices])
    train_loader = DataLoader(
        train, batch_size=512, shuffle=True, generator=generator, num_workers=0
    )
    val_loader = DataLoader(val, batch_size=512, shuffle=False, num_workers=0)
    torch.manual_seed(seed)
    model = ActionCopycat(action_dim=action_dim).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999))
    history = []
    for epoch in range(1, epochs + 1):
        model.train()
        losses = []
        for batch_x, batch_y in train_loader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)
            optimizer.zero_grad(set_to_none=True)
            loss = (model(batch_x) - batch_y).square().mean(dim=-1).mean()
            loss.backward()
            optimizer.step()
            losses.append(float(loss.detach().cpu()))
        model.eval()
        val_losses = []
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x = batch_x.to(device)
                batch_y = batch_y.to(device)
                loss = (model(batch_x) - batch_y).square().mean(dim=-1)
                val_losses.extend(loss.cpu().tolist())
        history.append(
            {
                "epoch": epoch,
                "train_loss": float(np.mean(losses)),
                "validation_loss": float(np.mean(val_losses)),
            }
        )
        print(json.dumps(history[-1], sort_keys=True), flush=True)
        if epoch % max(1, epochs // 3) == 0:
            for group in optimizer.param_groups:
                group["lr"] *= 0.1
    return model, history


def prediction_errors(model, inputs, targets, device):
    errors = []
    loader = DataLoader(
        TensorDataset(torch.from_numpy(inputs), torch.from_numpy(targets)),
        batch_size=4096,
        shuffle=False,
        num_workers=0,
    )
    model.eval()
    with torch.no_grad():
        for batch_x, batch_y in loader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)
            errors.append(
                (model(batch_x) - batch_y).square().mean(dim=-1).cpu().numpy()
            )
    return np.concatenate(errors)


def score_support(
    model,
    data_root,
    dataset_name,
    support,
    statistics,
    device,
    *,
    action_dim,
    binary_action_indices,
):
    trajectory_ids, offsets, anchors_out, ape_out = [], [0], [], []
    model.eval()
    with torch.no_grad():
        for trajectory_id, anchors in support.items():
            actions = q99_normalize(
                read_actions(data_root, dataset_name, trajectory_id, action_dim),
                statistics,
                binary_action_indices=binary_action_indices,
            )
            batch_x, batch_y = [], []
            for anchor in anchors:
                x, y = history_target(actions, anchor, history=9)
                batch_x.append(x)
                batch_y.append(y)
            x = torch.from_numpy(np.asarray(batch_x, dtype=np.float32)).to(device)
            y = torch.from_numpy(np.asarray(batch_y, dtype=np.float32)).to(device)
            ape = (model(x) - y).square().mean(dim=-1).cpu().numpy()
            trajectory_ids.append(trajectory_id)
            anchors_out.extend(anchors)
            ape_out.extend(ape.tolist())
            offsets.append(len(anchors_out))
    return {
        "trajectory_ids": np.asarray(trajectory_ids, dtype=np.int64),
        "offsets": np.asarray(offsets, dtype=np.int64),
        "anchors": np.asarray(anchors_out, dtype=np.int32),
        "ape": np.asarray(ape_out, dtype=np.float32),
    }


def verify_reference_overlap(reference, dataset_index, scored, weights):
    prefix = f"d{dataset_index}_"
    reference_trajectory_ids = reference[prefix + "trajectory_ids"]
    reference_offsets = reference[prefix + "offsets"]
    reference_anchors = reference[prefix + "anchors"]
    reference_ape = reference[prefix + "ape"]
    reference_weights = reference[prefix + "weights"]
    new_trajectory_positions = {
        int(trajectory_id): position
        for position, trajectory_id in enumerate(scored["trajectory_ids"])
    }
    checked = 0
    max_ape_difference = 0.0
    for reference_position, trajectory_id in enumerate(reference_trajectory_ids):
        new_position = new_trajectory_positions.get(int(trajectory_id))
        if new_position is None:
            raise RuntimeError(
                f"full cache is missing reference trajectory {int(trajectory_id)}"
            )
        reference_slice = slice(
            int(reference_offsets[reference_position]),
            int(reference_offsets[reference_position + 1]),
        )
        new_slice = slice(
            int(scored["offsets"][new_position]),
            int(scored["offsets"][new_position + 1]),
        )
        new_anchors = scored["anchors"][new_slice]
        positions = np.searchsorted(new_anchors, reference_anchors[reference_slice])
        if np.any(positions >= len(new_anchors)) or not np.array_equal(
            new_anchors[positions], reference_anchors[reference_slice]
        ):
            raise RuntimeError(
                f"full cache does not contain the reference anchor support for "
                f"trajectory {int(trajectory_id)}"
            )
        new_ape = scored["ape"][new_slice][positions]
        ape_difference = np.abs(new_ape - reference_ape[reference_slice])
        if len(ape_difference):
            max_ape_difference = max(max_ape_difference, float(ape_difference.max()))
        if not np.array_equal(
            weights[new_slice][positions], reference_weights[reference_slice]
        ):
            raise RuntimeError(
                f"full cache changed reference weights for trajectory "
                f"{int(trajectory_id)}"
            )
        checked += len(positions)
    if max_ape_difference > 1e-6:
        raise RuntimeError(
            f"full cache changed reference APE scores: max_abs_diff={max_ape_difference}"
        )
    return {
        "checked_anchor_count": checked,
        "weights_exact_match": True,
        "max_ape_abs_difference": max_ape_difference,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--statistics", type=Path, required=True)
    parser.add_argument(
        "--statistics-key-map",
        type=Path,
        help="Optional JSON mapping from manifest dataset name to statistics key.",
    )
    parser.add_argument(
        "--action-dim",
        type=int,
        choices=(7, 14),
        help="Expected action dimension; inferred from statistics when omitted.",
    )
    parser.add_argument(
        "--binary-action-indices",
        help=(
            "Comma-separated zero-based binary action indices. Defaults to index 6 "
            "for 7D Bridge and none for 14D RoboTwin absolute qpos."
        ),
    )
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--scope", choices=("full", "draw-range"), default="full")
    parser.add_argument("--draw-seed", type=int, default=117465)
    parser.add_argument("--draw-start", type=int, default=1_280_000)
    parser.add_argument("--draw-count", type=int, default=128)
    parser.add_argument("--fit-max-anchors", type=int)
    parser.add_argument("--epochs", type=int, default=300)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--reference-cache", type=Path)
    args = parser.parse_args()

    manifest = ValidAnchorManifest.load(args.manifest)
    statistics_payload = json.loads(args.statistics.read_text(encoding="utf-8"))
    statistics_key_map = (
        json.loads(args.statistics_key_map.read_text(encoding="utf-8"))
        if args.statistics_key_map is not None
        else {}
    )
    reference = None
    reference_metadata = None
    if args.reference_cache is not None:
        if args.scope != "full":
            raise ValueError("--reference-cache requires --scope full")
        reference = np.load(args.reference_cache, allow_pickle=False)
        reference_metadata = json.loads(str(reference["metadata_json"]))
        if int(reference_metadata["copycat_epochs"]) != args.epochs:
            raise RuntimeError("reference cache copycat epoch count does not match")
    scoring = requested_support(manifest, args)
    fitting = fitting_support(manifest, args.fit_max_anchors)
    arrays = {}
    dataset_reports = []
    for dataset_index, dataset in enumerate(manifest.datasets):
        dataset_name = dataset.name
        statistics_key = resolve_statistics_key(
            dataset_name, statistics_payload, statistics_key_map
        )
        action_statistics = statistics_payload[statistics_key]["action"]
        inferred_action_dim = len(action_statistics["q99"])
        action_dim = args.action_dim or inferred_action_dim
        if inferred_action_dim != action_dim:
            raise RuntimeError(
                f"statistics action dimension {inferred_action_dim} != {action_dim} "
                f"for {dataset_name}"
            )
        if args.binary_action_indices is None:
            binary_action_indices = (6,) if action_dim == 7 else ()
        else:
            binary_action_indices = tuple(
                int(value)
                for value in args.binary_action_indices.split(",")
                if value.strip()
            )
        inputs, targets = examples_for_support(
            args.data_root,
            dataset_name,
            fitting[dataset_name],
            action_statistics,
            action_dim=action_dim,
            binary_action_indices=binary_action_indices,
        )
        print(
            f"training copycat for {dataset_name}: {len(inputs)} anchors",
            flush=True,
        )
        model, train_history = train_copycat(
            inputs,
            targets,
            action_dim=action_dim,
            epochs=args.epochs,
            device=torch.device(args.device),
        )
        fitted_ape = prediction_errors(
            model, inputs, targets, torch.device(args.device)
        )
        scored = score_support(
            model,
            args.data_root,
            dataset_name,
            scoring[dataset_name],
            action_statistics,
            torch.device(args.device),
            action_dim=action_dim,
            binary_action_indices=binary_action_indices,
        )
        if reference_metadata is None:
            threshold_values = scored["ape"] if args.scope == "full" else fitted_ape
            threshold = float(np.quantile(threshold_values, 0.9))
        else:
            reference_report = reference_metadata["datasets"][dataset_index]
            if reference_report["name"] != dataset_name:
                raise RuntimeError("reference cache dataset order does not match")
            if int(reference_report["fit_anchor_count"]) != len(inputs):
                raise RuntimeError("reference cache fitting support does not match")
            threshold = float(reference_report["ape_threshold_90th_percentile"])
        weights = np.where(scored["ape"] >= threshold, 5.0, 1.0).astype(np.float32)
        overlap_report = None
        if reference is not None:
            overlap_report = verify_reference_overlap(
                reference, dataset_index, scored, weights
            )
        for key, value in scored.items():
            arrays[f"d{dataset_index}_{key}"] = value
        arrays[f"d{dataset_index}_weights"] = weights
        dataset_reports.append(
            {
                "name": dataset_name,
                "statistics_key": statistics_key,
                "action_dim": action_dim,
                "binary_action_indices": list(binary_action_indices),
                "fit_anchor_count": len(inputs),
                "scored_anchor_count": len(scored["anchors"]),
                "ape_threshold_90th_percentile": threshold,
                "upweighted_count": int((weights > 1).sum()),
                "copycat_training": train_history,
                "reference_overlap": overlap_report,
            }
        )

    metadata = {
        "schema_version": 1,
        "method": "Keyframe-Focused-adapted",
        "history_length": 9,
        "copycat_hidden_units": [300],
        "copycat_optimizer": {"name": "Adam", "lr": 1e-3, "betas": [0.9, 0.999]},
        "copycat_epochs": args.epochs,
        "ape_metric": "mean squared error over normalized action dimensions",
        "boundary_padding": "repeat earliest physical action",
        "threshold_scope": "within dataset source",
        "threshold_percentile": 90,
        "high_weight": 5.0,
        "low_weight": 1.0,
        "cache_scope": args.scope,
        "threshold_source": (
            "reference_cache" if args.reference_cache is not None else "computed"
        ),
        "dataset_names": [dataset.name for dataset in manifest.datasets],
        "datasets": dataset_reports,
    }
    arrays["metadata_json"] = np.asarray(json.dumps(metadata, sort_keys=True))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
    with temporary.open("wb") as destination:
        np.savez_compressed(destination, **arrays)
    temporary.replace(args.output)
    print(f"wrote {args.output}", flush=True)


if __name__ == "__main__":
    main()