"""Compute performance / score / perceptual codecs from aligned datasets.

Produces a single HDF5 file under ``data_root`` (one group per piece, segmented
to ``MAX_NOTE_LEN``) plus per-segment ``snote_ids`` ``.npy`` files.

Path configuration precedence: CLI flags > environment variables > config/paths.yaml.
"""

from __future__ import annotations

import argparse
import glob
import json
import logging
import os
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional

import h5py
import numpy as np
import numpy.lib.recfunctions as rfn
import pandas as pd
import partitura as pt
from omegaconf import OmegaConf
from tqdm import tqdm

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Path resolution
# ---------------------------------------------------------------------------

DEFAULT_PATHS_YAML = Path(__file__).parent / "config" / "paths.yaml"


@dataclass
class DatasetPaths:
    data_root: Path                 # where outputs (hdf5, snote_ids) are written
    datasets_root: Path             # parent of source corpora

    atepp_root: Path
    atepp_meta: Path
    asap_root: Path
    vienna_match_dir: Path
    vienna_musicxml_dir: Path
    vienna_cep_dir: Path


def load_paths(
    config_path: Optional[Path] = None,
    data_root: Optional[str] = None,
    datasets_root: Optional[str] = None,
) -> DatasetPaths:
    """Resolve dataset paths from yaml + environment + CLI overrides."""
    cfg = OmegaConf.load(config_path or DEFAULT_PATHS_YAML)

    resolved_data_root = Path(
        data_root
        or os.environ.get("DEXTER_DATA_ROOT")
        or cfg.data_root
    ).expanduser().resolve()

    resolved_datasets_root = Path(
        datasets_root
        or os.environ.get("DEXTER_DATASETS_ROOT")
        or cfg.datasets_root
    ).expanduser().resolve()

    return DatasetPaths(
        data_root=resolved_data_root,
        datasets_root=resolved_datasets_root,
        atepp_root=resolved_datasets_root / cfg.datasets.ATEPP.root,
        atepp_meta=resolved_datasets_root / cfg.datasets.ATEPP.metadata,
        asap_root=resolved_datasets_root / cfg.datasets.ASAP.root,
        vienna_match_dir=resolved_datasets_root / cfg.datasets.VIENNA422.match_dir,
        vienna_musicxml_dir=resolved_datasets_root / cfg.datasets.VIENNA422.musicxml_dir,
        vienna_cep_dir=resolved_datasets_root / cfg.datasets.VIENNA422.cep_dir,
    )


# ---------------------------------------------------------------------------
# Per-dataset path discovery
# ---------------------------------------------------------------------------

@dataclass
class PieceTuple:
    score_path: Path
    performance_path: Optional[Path]   # None for VIENNA422 (alignment carries the perf)
    alignment_path: Path
    cep_path: Path
    piece_name: str


def discover_vienna(paths: DatasetPaths) -> list[PieceTuple]:
    alignment_paths = sorted(glob.glob(str(paths.vienna_match_dir / "*[!e].match")))
    out = []
    for a in alignment_paths:
        a_p = Path(a)
        stem = a_p.name[:-10]               # strip "_pNN.match"
        out.append(PieceTuple(
            score_path=paths.vienna_musicxml_dir / f"{stem}.musicxml",
            performance_path=None,
            alignment_path=a_p,
            cep_path=paths.vienna_cep_dir / f"{a_p.stem}.csv",
            piece_name=stem,
        ))
    return out


def discover_asap(paths: DatasetPaths) -> list[PieceTuple]:
    perf_paths = sorted(glob.glob(str(paths.asap_root / "**/*[!e].mid"), recursive=True))
    out = []
    for p in perf_paths:
        p_p = Path(p)
        out.append(PieceTuple(
            score_path=p_p.parent / "xml_score.musicxml",
            performance_path=p_p,
            alignment_path=p_p.parent / f"{p_p.stem}_note_alignments" / "note_alignment.tsv",
            cep_path=p_p.parent / f"{p_p.stem}_cep_features.csv",
            piece_name=str(p_p.relative_to(paths.asap_root).parent).replace("/", "_"),
        ))
    return out


def discover_atepp(paths: DatasetPaths) -> list[PieceTuple]:
    """ATEPP layout: <piece_dir>/<seg_id>/align.csv with perf at <piece_dir>/<seg_id>.mid."""
    alignment_paths = sorted(glob.glob(str(paths.atepp_root / "**/[!z]*n.csv"), recursive=True))
    out = []
    for a in alignment_paths:
        a_p = Path(a)
        seg_id = a_p.parent.name
        piece_dir = a_p.parent.parent
        perf = piece_dir / f"{seg_id}.mid"
        score_candidates = glob.glob(str(piece_dir / "*.*l"))   # .xml, .mxl, .musicxml
        if not score_candidates:
            continue
        out.append(PieceTuple(
            score_path=Path(score_candidates[0]),
            performance_path=perf,
            alignment_path=a_p,
            cep_path=piece_dir / f"{seg_id}_cep_features.csv",
            piece_name=seg_id,
        ))
    return out


DISCOVERERS = {
    "VIENNA422": discover_vienna,
    "ASAP": discover_asap,
    "ATEPP": discover_atepp,
}


def get_atepp_overlap(paths: DatasetPaths) -> set[str]:
    """ATEPP score folders with multiple performances (mixup-eligible)."""
    meta = pd.read_csv(paths.atepp_meta)
    score_groups = meta.groupby(["score_path"])["midi_path"].count()
    return {
        str(paths.atepp_root / Path(score_path).parent)
        for score_path in score_groups.index
    }


# ---------------------------------------------------------------------------
# Codec extraction (per piece)
# ---------------------------------------------------------------------------

def get_codecs(score_path, alignment_path, cep_path, performance_path=None, score=None):
    """Encode a single (score, performance, alignment) triple into p_codec / c_codec."""
    score_path = str(score_path)
    alignment_path = str(alignment_path)
    cep_path = str(cep_path) if cep_path is not None else None

    if alignment_path.endswith(".match"):
        performance, alignment = pt.load_match(alignment_path)
    elif alignment_path.endswith(".tsv"):
        alignment = pt.io.importparangonada.load_alignment_from_ASAP(alignment_path)
    elif alignment_path.endswith(".csv"):                     # ATEPP / parangonar
        alignment = pt.io.importparangonada.load_parangonada_alignment(alignment_path)
        if score is None:
            score = pt.load_musicxml(score_path, force_note_ids="keep")
    else:
        raise ValueError(f"Unknown alignment format: {alignment_path}")

    if score is None:
        score = pt.load_musicxml(score_path)

    # if the alignment has unfolded score IDs but the score doesn't, unfold the score
    if (("score_id" in alignment[0])
            and ("-" in alignment[0]["score_id"])
            and ("-" not in score.note_array(include_divs_per_quarter=False)["id"][0])):
        score = pt.score.unfold_part_maximal(pt.score.merge_parts(score.parts))

    if performance_path is not None:
        performance = pt.load_performance(str(performance_path))

    parameters, snote_ids, _, m_score = pt.musicanalysis.encode_performance(
        score, performance, alignment,
    )
    p_codec = rfn.structured_to_unstructured(parameters)

    if cep_path and os.path.exists(cep_path):
        cep_feats = pd.read_csv(cep_path)
        c_codec = rfn.structured_to_unstructured(get_cep_codec(cep_feats, m_score))
    else:
        if cep_path:
            logger.info("No cep_features for %s", alignment_path)
        c_codec = np.zeros((len(p_codec), 7))

    return p_codec, c_codec, snote_ids, score


def get_cep_codec(cep_feats: pd.DataFrame, m_score) -> np.ndarray:
    """Align perceptual features with the performance, return structured array."""
    rows = []
    for note in m_score:
        window = cep_feats[cep_feats["frame_start_time"] <= note["p_onset"]]
        if len(window):
            rows.append(window.iloc[-1])
        else:
            rows.append(cep_feats.iloc[-1])

    c_codec_df = pd.DataFrame(rows).drop(columns=["frame_start_time"])
    records = c_codec_df.to_records(index=False)
    return np.array(records, dtype=records.dtype.descr)


# ---------------------------------------------------------------------------
# Main processing loop
# ---------------------------------------------------------------------------

def hdf5_path(paths: DatasetPaths, max_note_len: int, mixup: bool) -> Path:
    suffix = "_mixup" if mixup else ""
    return paths.data_root / f"codec_N={max_note_len}{suffix}.hdf5"


def process_dataset_codec(
    paths: DatasetPaths,
    max_note_len: int,
    mix_up: bool = False,
    datasets: Iterable[str] = ("ATEPP", "ASAP", "VIENNA422"),
    test_split_p: float = 0.15,
    rng_seed: Optional[int] = None,
) -> Path:
    """Compute and persist codecs for the requested datasets.

    Returns the path to the written HDF5 file.
    Per-piece errors are logged to ``data_root / errors.jsonl``.
    """
    rng = np.random.default_rng(rng_seed)
    paths.data_root.mkdir(parents=True, exist_ok=True)
    snote_dir = paths.data_root / f"snote_ids/N={max_note_len}"
    snote_dir.mkdir(parents=True, exist_ok=True)

    out_path = hdf5_path(paths, max_note_len, mix_up)
    errors_path = paths.data_root / "errors.jsonl"

    atepp_overlap = get_atepp_overlap(paths) if "ATEPP" in datasets else set()

    with h5py.File(out_path, "a") as hdf5_file, errors_path.open("a") as err_log:
        for dataset in datasets:
            pieces = DISCOVERERS[dataset](paths)
            logger.info("[%s] discovered %d pieces", dataset, len(pieces))

            same_score_p, mixup_p = [], []
            same_score_c, mixup_c = [], []
            prev_score_path = None
            split = "train"

            for j, piece in enumerate(tqdm(pieces, desc=dataset)):
                group_key = f"{dataset}_{j}"
                if group_key in hdf5_file:
                    continue
                if not (piece.score_path.exists() and piece.alignment_path.exists()):
                    err_log.write(json.dumps({
                        "dataset": dataset, "idx": j, "reason": "missing_files",
                        "score": str(piece.score_path), "alignment": str(piece.alignment_path),
                    }) + "\n")
                    continue

                # ATEPP-only: drop alignments with <50% match
                if dataset == "ATEPP":
                    try:
                        alignment = pt.io.importparangonada.load_parangonada_alignment(str(piece.alignment_path))
                    except Exception as e:
                        err_log.write(json.dumps({
                            "dataset": dataset, "idx": j, "reason": "alignment_load_error",
                            "alignment": str(piece.alignment_path), "error": str(e),
                        }) + "\n")
                        continue
                    counts = {label: 0 for label in ("match", "insertion", "deletion")}
                    for a in alignment:
                        if a["label"] in counts:
                            counts[a["label"]] += 1
                    total = sum(counts.values())
                    if total == 0 or counts["match"] / total < 0.5:
                        err_log.write(json.dumps({
                            "dataset": dataset, "idx": j, "reason": "low_alignment_match_ratio",
                            "ratio": counts["match"] / total if total else 0.0,
                        }) + "\n")
                        continue

                reuse_score = (prev_score_path == piece.score_path)
                try:
                    p_codec, c_codec, snote_ids, score = get_codecs(
                        piece.score_path, piece.alignment_path, piece.cep_path,
                        performance_path=piece.performance_path,
                        score=(score if reuse_score else None),
                    )
                except Exception as e:
                    err_log.write(json.dumps({
                        "dataset": dataset, "idx": j, "reason": "codec_extraction_error",
                        "score": str(piece.score_path), "error": str(e),
                    }) + "\n")
                    continue

                # tempo sanity check
                if 60 / p_codec[:, 0].mean() > 200:
                    hdf5_file.create_group(group_key).create_dataset("error", data="tempo_outlier")
                    continue

                if not reuse_score:
                    same_score_p, mixup_p = [], []
                    same_score_c, mixup_c = [], []
                    split = "train" if rng.random() > test_split_p else "test"
                else:
                    eligible = (
                        dataset == "VIENNA422"
                        or (dataset == "ATEPP" and str(piece.score_path.parent) in atepp_overlap)
                    )
                    if mix_up and eligible:
                        mixup_p = [np.mean([p_codec, ssp], axis=0) for ssp in same_score_p]
                        mixup_c = [np.mean([c_codec, ssc], axis=0) for ssc in same_score_c]
                        same_score_p.append(p_codec)
                        same_score_c.append(c_codec)

                sna = score.note_array()
                sna = sna[np.in1d(sna["id"], snote_ids)]
                s_codec = rfn.structured_to_unstructured(
                    sna[["onset_div", "duration_div", "pitch", "voice"]]
                )

                if not (len(p_codec) == len(s_codec) == len(snote_ids)):
                    hdf5_file.create_group(group_key).create_dataset("error", data="length_mismatch")
                    err_log.write(json.dumps({
                        "dataset": dataset, "idx": j, "reason": "length_mismatch",
                        "p": len(p_codec), "s": len(s_codec), "snote_ids": len(snote_ids),
                    }) + "\n")
                    continue

                data = _segment_piece(
                    snote_dir=snote_dir,
                    dataset=dataset,
                    piece_name=piece.piece_name,
                    score_path=piece.score_path,
                    p_codecs=[p_codec, *mixup_p],
                    c_codecs=[c_codec, *mixup_c],
                    s_codec=s_codec,
                    snote_ids=snote_ids,
                    split=split,
                    max_note_len=max_note_len,
                )

                grp = hdf5_file.create_group(group_key)
                for key in ("p_codec", "s_codec", "c_codec", "snote_id_path",
                            "score_path", "piece_name"):
                    grp.create_dataset(key, data=data[key])
                grp.create_dataset(
                    "split", data=data["split"], dtype=h5py.string_dtype(encoding="utf-8")
                )

                prev_score_path = piece.score_path

    logger.info("Wrote %s", out_path)
    return out_path


def _segment_piece(
    *, snote_dir, dataset, piece_name, score_path,
    p_codecs, c_codecs, s_codec, snote_ids, split, max_note_len,
):
    """Segment a piece (and its mixup variants) into fixed-length windows.

    Pads short tail segments to ``max_note_len`` with zeros.
    """
    data = defaultdict(list)
    for variant_idx, (p_codec, c_codec) in enumerate(zip(p_codecs, c_codecs)):
        variant_name = piece_name + ("_mixup" if variant_idx > 0 else "")
        save_id_prefix = snote_dir / f"{dataset}_{piece_name}"

        for seg_idx, start in enumerate(range(0, len(p_codec), max_note_len)):
            end = start + max_note_len
            seg_p = p_codec[start:end]
            seg_s = s_codec[start:end]
            seg_c = c_codec[start:end]
            seg_ids = snote_ids[start:end]
            if len(seg_ids) == 0:
                continue

            if len(seg_p) < max_note_len:
                pad = max_note_len - len(seg_p)
                seg_p = np.pad(seg_p, ((0, pad), (0, 0)))
                seg_s = np.pad(seg_s, ((0, pad), (0, 0)))
                seg_c = np.pad(seg_c, ((0, pad), (0, 0)))

            seg_id_path = f"{save_id_prefix}_seg{seg_idx}.npy"
            np.save(seg_id_path, seg_ids)

            data["p_codec"].append(seg_p)
            data["s_codec"].append(seg_s)
            data["c_codec"].append(seg_c)
            data["snote_id_path"].append(seg_id_path)
            data["score_path"].append(str(score_path))
            data["piece_name"].append(variant_name)
            data["split"].append(split)
    return data


# ---------------------------------------------------------------------------
# HDF5 → in-memory loader (consumed by train.py)
# ---------------------------------------------------------------------------

def _decode(x) -> str:
    """h5py returns numpy bytes for string datasets; decode to plain str."""
    return x.decode() if hasattr(x, "decode") else str(x)


def _apply_path_remap(path: str, path_remap: Optional[dict]) -> str:
    """Prefix-substitute a path, first match wins. Pass-through if no remap matches."""
    if not path_remap:
        return path
    for old, new in path_remap.items():
        if path.startswith(old):
            return new + path[len(old):]
    return path


def load_data_from_hdf5(hdf5_path, path_remap: Optional[dict] = None):
    """Load codec dicts from HDF5, optionally rewriting absolute paths.

    ``path_remap`` is a dict of ``{old_prefix: new_prefix}``; both ``score_path``
    and ``snote_id_path`` are rewritten when they start with one of the keys.
    Use this to consume an HDF5 produced on a different machine — score paths
    embedded by ``prepare_data.py`` are absolute, so cross-host portability needs
    a runtime override.
    """
    train_data, test_data = [], []
    logger.info("Loading data from %s", hdf5_path)
    if path_remap:
        logger.info("path_remap: %s", dict(path_remap))

    with h5py.File(hdf5_path, "r") as hdf5_file:
        for group_name in tqdm(hdf5_file):
            group = hdf5_file[group_name]
            if "error" in group:
                continue

            p_codec = np.array(group["p_codec"])
            s_codec = np.array(group["s_codec"])
            c_codec = np.array(group["c_codec"])
            snote_id_path = np.array(group["snote_id_path"])
            score_path = np.array(group["score_path"])
            piece_name = np.array(group["piece_name"])
            split = np.array(group["split"])

            piece_data = [{
                "p_codec": p_codec[i],
                "s_codec": s_codec[i],
                "c_codec": c_codec[i],
                "snote_id_path": _apply_path_remap(_decode(snote_id_path[i]), path_remap),
                "score_path": _apply_path_remap(_decode(score_path[i]), path_remap),
                "piece_name": _decode(piece_name[i]),
                "split": _decode(split[i]),
            } for i in range(len(p_codec))]

            if piece_data[0]["split"] == "train":
                train_data.extend(piece_data)
            else:
                test_data.extend(pd_ for pd_ in piece_data if "mixup" not in pd_["piece_name"])

    return train_data, test_data


# ---------------------------------------------------------------------------
# Transfer pairing (segment-level pairs of real performances of the same piece)
# ---------------------------------------------------------------------------

def make_transfer_pair(codec_data, paths: DatasetPaths, K=50000, N=200):
    """Generate K pairs of segments (same piece + same segment idx, different performance).

    Cached at ``data_root/codec_N={N}_mixup_paired_K={K}.npy`` and ``..._unpaired_K={K}.npy``.
    """
    paired_path = paths.data_root / f"codec_N={N}_mixup_paired_K={K}.npy"
    unpaired_path = paths.data_root / f"codec_N={N}_mixup_unpaired_K={K}.npy"
    if paired_path.exists() and unpaired_path.exists():
        return (
            np.load(paired_path, allow_pickle=True),
            np.load(unpaired_path, allow_pickle=True),
        )

    codec_data = np.array(codec_data)
    real_only = codec_data[[i for i, x in enumerate(codec_data) if "mu" not in x["piece_name"]]]
    np.random.shuffle(real_only)
    unpaired = codec_data
    transfer_pairs = []

    for cd in tqdm(real_only):
        if len(transfer_pairs) > K:
            break
        seg_id = cd["snote_id_path"].split("seg")[-1].split(".")[0]
        mask = [
            (x["score_path"] == cd["score_path"]
             and f"seg{seg_id}." in x["snote_id_path"]
             and (x["p_codec"] != cd["p_codec"]).any())
            for x in real_only
        ]
        same_piece = real_only[mask]
        if len(same_piece):
            for scd in same_piece:
                transfer_pairs.extend([cd, scd])
            unpaired = unpaired[[x["score_path"] != cd["score_path"] for x in unpaired]]

    transfer_pairs = np.array(transfer_pairs).reshape(2, -1, order="F")
    np.random.shuffle(transfer_pairs.T)
    transfer_pairs = transfer_pairs.ravel(order="F")

    np.save(paired_path, transfer_pairs)
    np.save(unpaired_path, unpaired)
    return transfer_pairs, unpaired


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="DExter codec preprocessing")
    p.add_argument("--compute_codec", action="store_true",
                   help="compute and persist p/s/c codecs to HDF5")
    p.add_argument("--pairing", action="store_true",
                   help="build transfer pairs from a precomputed codec HDF5")
    p.add_argument("--MAX_NOTE_LEN", type=int, default=200,
                   help="segment length in notes")
    p.add_argument("--K", type=int, default=2000,
                   help="number of transfer pairs to draw")
    p.add_argument("--mixup", action="store_true", default=True,
                   help="enable mixup augmentation across same-score performances")
    p.add_argument("--no-mixup", dest="mixup", action="store_false")
    p.add_argument("--datasets", nargs="+", default=["ATEPP", "ASAP", "VIENNA422"],
                   choices=list(DISCOVERERS.keys()))
    p.add_argument("--data_root", type=str, default=None,
                   help="output directory; overrides $DEXTER_DATA_ROOT and config/paths.yaml")
    p.add_argument("--datasets_root", type=str, default=None,
                   help="parent of source corpora; overrides $DEXTER_DATASETS_ROOT")
    p.add_argument("--paths_config", type=str, default=None,
                   help="path to a paths.yaml override")
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--log_level", default="INFO")
    return p


def main():
    args = _build_argparser().parse_args()
    logging.basicConfig(level=args.log_level, format="%(asctime)s [%(levelname)s] %(message)s")

    paths = load_paths(
        config_path=Path(args.paths_config) if args.paths_config else None,
        data_root=args.data_root,
        datasets_root=args.datasets_root,
    )
    logger.info("data_root=%s  datasets_root=%s", paths.data_root, paths.datasets_root)

    if args.compute_codec:
        process_dataset_codec(
            paths,
            max_note_len=args.MAX_NOTE_LEN,
            mix_up=args.mixup,
            datasets=args.datasets,
            rng_seed=args.seed,
        )
    elif args.pairing:
        codec_path = paths.data_root / f"codec_N={args.MAX_NOTE_LEN}_mixup_test.npy"
        codec_data = np.load(codec_path, allow_pickle=True)
        make_transfer_pair(codec_data, paths, K=args.K, N=args.MAX_NOTE_LEN)
    else:
        _build_argparser().print_help()


if __name__ == "__main__":
    main()