"""Real RoboTwin joint/gripper state IO and train-anchor-only normalization."""
from __future__ import annotations

from contextlib import ExitStack
import json
from pathlib import Path
from tempfile import TemporaryDirectory

import numpy as np
import torch

STATE_NORM_FILENAME = "state_norm.json"
STATE_LAYOUT = [f"{arm}.{name}" for arm in ("left", "right")
                for name in (*[f"joint_{i}" for i in range(6)], "gripper")]
JOINT_INDICES = [i for i in range(14) if i not in (6, 13)]


def read_robot_states(episode_path, num_frames=None):
    path = Path(episode_path) / "robot_state.npy"
    if not path.is_file():
        raise ValueError(f"V2 requires real current robot state: missing {path}")
    state = np.load(path, mmap_mode="r", allow_pickle=False)
    if state.ndim != 2 or state.shape[1] != 14 or (num_frames is not None and len(state) != num_frames):
        raise ValueError(f"Expected frame-aligned robot_state.npy [T,14]: {path}, got {state.shape}")
    return state


def validate_robot_states(values):
    values = np.asarray(values)
    if values.ndim != 2 or values.shape[1] != 14 or not np.isfinite(values).all():
        raise ValueError("Current real robot states must be finite [N,14]")
    if np.any((values[:, [6, 13]] < 0) | (values[:, [6, 13]] > 1)):
        raise ValueError("Current real robot state grippers must be in [0,1]")


def state_parameters(q01, q99, minimum, maximum):
    q01, q99, minimum, maximum = [np.asarray(x, dtype=np.float64) for x in (q01, q99, minimum, maximum)]
    if any(x.shape != (1, 12) or not np.isfinite(x).all() for x in (q01, q99, minimum, maximum)):
        raise ValueError("State joint statistics must be finite [1,12]")
    if np.any((minimum > q01) | (q01 > q99) | (q99 > maximum)):
        raise ValueError("Invalid state quantile bounds")
    fallback = q99 - q01 < 1e-4
    low, high = np.where(fallback, minimum, q01), np.where(fallback, maximum, q99)
    offset, scale = np.full((1, 14), 0.5), np.full((1, 14), 0.5)
    offset[:, JOINT_INDICES] = (low + high) / 2
    scale[:, JOINT_INDICES] = np.maximum(high - low, 1e-3) / 2
    return offset, scale


def validate_state_norm(record):
    expected = {"version": 1, "layout": STATE_LAYOUT, "statistics_scope": "train_valid_current_frames",
                "joint_indices": JOINT_INDICES, "minimum_span": 1e-3, "fallback_span": 1e-4}
    for key, value in expected.items():
        if record.get(key) != value:
            raise ValueError(f"Invalid state normalization {key}")
    if not isinstance(record.get("count"), int) or record["count"] < 1:
        raise ValueError("State statistics require valid training frames")
    offset, scale = state_parameters(*(record[key] for key in ("q01", "q99", "min", "max")))
    for name, expected in (("offset", offset), ("scale", scale)):
        actual = np.asarray(record[name], dtype=np.float32)
        if actual.shape != (1, 14) or not np.isfinite(actual).all() or not np.array_equal(actual, expected.astype(np.float32)):
            raise ValueError(f"State normalization {name} disagrees with quantile/gripper rules")
    return record


def compute_state_norm(mixture, scratch_dir):
    """Exact quantiles per joint, bounded RAM; each eligible current frame once."""
    count, sources = 0, []
    minimum, maximum = np.full(12, np.inf), np.full(12, -np.inf)
    with TemporaryDirectory(prefix=".state-norm-", dir=scratch_dir) as tmp:
        paths = [Path(tmp) / f"joint-{i}.bin" for i in range(12)]
        with ExitStack() as stack:
            handles = [stack.enter_context(path.open("wb")) for path in paths]
            for source in mixture.sources:
                if source.weight <= 0:
                    continue
                if source.dataset.split != "train":
                    raise ValueError("State normalization may only use the training split")
                source_count = 0
                for values in source.dataset.iter_current_robot_states():
                    validate_robot_states(values)
                    joints = np.asarray(values, dtype=np.float32)[:, JOINT_INDICES]
                    if not len(joints):
                        continue
                    minimum, maximum = np.minimum(minimum, joints.min(0)), np.maximum(maximum, joints.max(0))
                    count += len(joints)
                    source_count += len(joints)
                    for column, handle in enumerate(handles):
                        np.ascontiguousarray(joints[:, column]).tofile(handle)
                sources.append({"name": source.name, "count": source_count, "selection": source.dataset.manifest()})
        if not count:
            raise ValueError("No valid training current frames for state normalization")
        ranks = np.array([0.01, 0.99]) * (count - 1)
        lower, upper = np.floor(ranks).astype(np.int64), np.ceil(ranks).astype(np.int64)
        quantiles = np.empty((2, 12))
        for column, path in enumerate(paths):
            values = np.memmap(path, mode="r+", dtype=np.float32, shape=(count,))
            values.partition(np.unique(np.r_[lower, upper]))
            lo, hi = values[lower].astype(np.float64), values[upper].astype(np.float64)
            quantiles[:, column] = lo + (hi - lo) * (ranks - lower)
            del values
    offset, scale = state_parameters(quantiles[:1], quantiles[1:], minimum[None], maximum[None])
    return validate_state_norm({"version": 1, "layout": STATE_LAYOUT, "joint_indices": JOINT_INDICES,
        "statistics_scope": "train_valid_current_frames", "minimum_span": 1e-3, "fallback_span": 1e-4,
        "count": count, "sources": sources, "q01": quantiles[:1].tolist(), "q99": quantiles[1:].tolist(),
        "min": minimum[None].tolist(), "max": maximum[None].tolist(),
        "offset": offset.tolist(), "scale": scale.tolist()})


def save_state_norm(directory, record):
    validate_state_norm(record)
    path = Path(directory) / STATE_NORM_FILENAME
    temporary = path.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(record, indent=2, allow_nan=False), encoding="utf-8")
    temporary.replace(path)


def load_checkpoint_state_norm(checkpoint):
    path = Path(checkpoint).expanduser()
    directory = path if path.is_dir() else path.parent
    record = json.loads((directory / STATE_NORM_FILENAME).read_text())
    return validate_state_norm(record)


def assert_state_norm_matches(policy, record):
    validate_state_norm(record)
    for name in ("offset", "scale"):
        actual = getattr(policy, f"state_{name}").detach().float().cpu()
        if not torch.equal(actual, torch.tensor(record[name], dtype=torch.float32)):
            raise ValueError(f"Checkpoint state_{name} differs from {STATE_NORM_FILENAME}")


def initialize_state_normalization(policy, dataset, config, accelerator):
    from accelerate.utils import broadcast_object_list

    result = [None]
    if accelerator.is_main_process:
        try:
            root = Path(config["output_dir"])
            root.mkdir(parents=True, exist_ok=True)
            record = load_checkpoint_state_norm(config["resume"]) if config.get("resume") else compute_state_norm(dataset, root)
            save_state_norm(root, record)
            result[0] = {"record": record}
        except Exception as exc:
            result[0] = {"error": f"{type(exc).__name__}: {exc}"}
    broadcast_object_list(result)
    if "error" in result[0]:
        raise ValueError(result[0]["error"])
    record = result[0]["record"]
    if config.get("resume"):
        assert_state_norm_matches(policy, record)
    else:
        policy.set_state_stats(record["offset"], record["scale"])
    return record