
"""
ncuhep.muography.data_processing_gui.data_processing

Core data-processing utilities used by the 4444 YMS GUI.

This module is intentionally independent from any GUI framework so it can be
used directly from scripts or notebooks.

Main entry points
-----------------
- parse_run_filename()
- filter_files()
- build_layer_geometry()
- compute_layer_pixel_mapping()
- compute_hitmaps_parallel()
- compute_flux_from_tracks()
"""

from __future__ import annotations

import logging
import os
import re
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import date, datetime
from pathlib import Path
from typing import Callable, Iterable, List, Optional, Sequence, Tuple

import numpy as np

from ncuhep.muography.classes import PlaneDetector, MuTxtFormat, AnalysisConfig
from ncuhep.units import Counts, GeometricFactor, Time

logger = logging.getLogger("4444_gui")


# ---------------------------------------------------------------------------
# Filename parsing & filtering
# ---------------------------------------------------------------------------


_RUN_RE = re.compile(
    r"(?P<date>\d{8})_Run_(?P<run>\d+)_(?P<sub>\d+)_Mu\.txt$", re.IGNORECASE
)


def parse_run_filename(name: str) -> Optional[Tuple[date, int, int]]:
    """
    Parse a run filename like '20251122_Run_3_0_Mu.txt'.

    Returns
    -------
    (day, run_id, sub_id) or None if not matched.
    """
    m = _RUN_RE.match(name)
    if not m:
        return None
    day = datetime.strptime(m.group("date"), "%Y%m%d").date()
    run_id = int(m.group("run"))
    sub_id = int(m.group("sub"))
    return day, run_id, sub_id


@dataclass
class FileFilterResult:
    selected: List[str]
    description: str


def filter_files(
    filenames: Sequence[str],
    mode: str,
    date_from: Optional[date] = None,
    date_to: Optional[date] = None,
    run_min: Optional[int] = None,
    run_max: Optional[int] = None,
) -> FileFilterResult:
    """
    Filter *_Mu.txt filenames according to GUI settings.

    Parameters
    ----------
    filenames:
        All candidate filenames (basename only).
    mode:
        'date' or 'run'.
    date_from, date_to:
        Used when mode == 'date'.
    run_min, run_max:
        Used when mode == 'run'.

    Returns
    -------
    FileFilterResult
    """
    selected: List[str] = []

    if mode == "date":
        if date_from is None:
            date_from = date.min
        if date_to is None:
            date_to = date.max
        for fn in filenames:
            parsed = parse_run_filename(fn)
            if parsed is None:
                continue
            d, _, _ = parsed
            if date_from <= d <= date_to:
                selected.append(fn)
        desc = f"Date {date_from:%Y-%m-%d} \u2192 {date_to:%Y-%m-%d}"
    elif mode == "run":
        if run_min is None:
            run_min = 0
        if run_max is None:
            run_max = 999999
        for fn in filenames:
            parsed = parse_run_filename(fn)
            if parsed is None:
                continue
            _, run_id, _ = parsed
            if run_min <= run_id <= run_max:
                selected.append(fn)
        desc = f"Run {run_min} \u2192 {run_max}"
    else:
        raise ValueError(f"Unknown filter mode: {mode}")

    return FileFilterResult(selected=sorted(selected), description=desc)


# ---------------------------------------------------------------------------
# Detector geometry helpers
# ---------------------------------------------------------------------------


def build_layer_geometry(det: PlaneDetector) -> Tuple[int, int, int, int, int]:
    """
    Compute basic geometry derived from a PlaneDetector.

    Returns
    -------
    (layer_count, nx, ny, ncell_layer, ncell_total)
      where nx, ny are channels along x/y for each layer.
    """
    # For the current 4444 detector we assume each layer is 8x8.
    # If in future config is more general this can be extended.
    layer_ids = np.asarray(det.layer_id, dtype=int)
    layer_count = layer_ids.size
    nx = int(det.layer_nx) if hasattr(det, "layer_nx") else 8
    ny = int(det.layer_ny) if hasattr(det, "layer_ny") else 8

    ncell_layer = nx * ny
    ncell_total = layer_count * ncell_layer
    return layer_count, nx, ny, ncell_layer, ncell_total


def compute_layer_pixel_mapping(det: PlaneDetector) -> np.ndarray:
    """
    Build a (layer_count, ny, nx) array giving the global UNIQUEID mapping.

    The convention follows the existing analysis script:
      layer 0: ids 0..63
      layer 1: ids 64..127
      layer 2: ids 128..191
      layer 3: ids 192..255
    """
    layer_count, nx, ny, ncell_layer, _ = build_layer_geometry(det)

    layers = np.full((layer_count, ny, nx), -1, dtype=int)

    for layer in range(layer_count):
        base = layer * ncell_layer
        for iy in range(ny):
            for ix in range(nx):
                uid = base + iy * nx + ix
                layers[layer, iy, ix] = uid

    return layers


# ---------------------------------------------------------------------------
# Parallel hitmap / track processing
# ---------------------------------------------------------------------------

# Globals used by worker processes (initialised via _init_worker)
_det: Optional[PlaneDetector] = None
_mutxt: Optional[MuTxtFormat] = None
_ana: Optional[AnalysisConfig] = None
_data_folder: Optional[Path] = None


def _init_worker(
    detector: PlaneDetector,
    mutxt: MuTxtFormat,
    analysis: AnalysisConfig,
    data_folder: str,
):
    global _det, _mutxt, _ana, _data_folder
    _det = detector
    _mutxt = mutxt
    _ana = analysis
    _data_folder = Path(data_folder)
    logger.info(
        "[worker init] detector=%s mutxt=%s analysis=%s data_folder=%s",
        detector,
        mutxt,
        analysis,
        data_folder,
    )


def _uids_to_hist_global(uids: np.ndarray, ncell_total: int) -> np.ndarray:
    """
    Convert an array of UNIQUEID values (0..N-1) into a flat histogram
    of length ncell_total.
    """
    h = np.bincount(uids.astype(int).ravel(), minlength=ncell_total)
    if h.size > ncell_total:
        h = h[:ncell_total]
    return h.astype(np.int64)


def _process_one(filename: str):
    """
    Worker function: read one *_Mu.txt file, build histograms and
    collect track angles + live time.
    """
    from ncuhep.muography import parser, identifier, tracker

    assert _det is not None and _mutxt is not None and _ana is not None
    assert _data_folder is not None

    src = _data_folder / filename
    filesize_mb = -1.0
    try:
        filesize_mb = os.path.getsize(src) / (1024.0 * 1024.0)
    except OSError:
        pass

    try:
        layer_count, nx, ny, ncell_layer, ncell_total = build_layer_geometry(_det)

        logger.info("[worker] processing file=%s", filename)

        events, timeelapsed, uids_hits = parser(
            str(_data_folder), filename, _mutxt, _det, _ana, True
        )

        uids_events = events["UNIQUEID"]

        events_id = identifier(events, _det)
        tracks = tracker(events_id, _det)

        if len(tracks) > 0:
            uids_tracks = np.concatenate([t["UNIQUEID"] for t in tracks])
            tracks_arr = np.asarray([t["TRACK"] for t in tracks], dtype=float)
            # TRACK = (dx, dy, tx, ty)
            tx = tracks_arr[:, 2].astype(float)
            ty = tracks_arr[:, 3].astype(float)
        else:
            uids_tracks = np.empty(0, dtype=int)
            tx = np.empty(0, dtype=float)
            ty = np.empty(0, dtype=float)

        if hasattr(timeelapsed, "s"):
            live_time_s = float(timeelapsed.s)
        else:
            live_time_s = float(timeelapsed)

        h_hits = _uids_to_hist_global(uids_hits, ncell_total)
        h_events = _uids_to_hist_global(uids_events, ncell_total)
        h_tracks = _uids_to_hist_global(uids_tracks, ncell_total)

        n_events = int(len(events))
        n_tracks = int(len(tracks))

        logger.info(
            "[worker] done file=%s size=%.2f MB events=%d tracks=%d time=%.1fs",
            filename,
            filesize_mb,
            n_events,
            n_tracks,
            live_time_s,
        )

        return (
            filename,
            True,
            "",
            filesize_mb,
            h_hits,
            h_events,
            h_tracks,
            tx,
            ty,
            live_time_s,
            n_events,
            n_tracks,
        )
    except Exception as exc:  # noqa: BLE001
        logger.exception("[worker] exception in file=%s", filename)
        return (
            filename,
            False,
            str(exc),
            filesize_mb,
            None,
            None,
            None,
            None,
            None,
            0.0,
            0,
            0,
        )


@dataclass
class HitmapResult:
    hits_per_layer: List[np.ndarray]
    events_per_layer: List[np.ndarray]
    tracks_per_layer: List[np.ndarray]
    detector: PlaneDetector
    summary: str
    failures: List[Tuple[str, str, float]]
    tx_all: Optional[np.ndarray]
    ty_all: Optional[np.ndarray]
    total_live_time_s: float
    total_events: int
    total_tracks: int


def compute_hitmaps_parallel(
    data_folder: str,
    config_folder: str,
    filenames: Sequence[str],
    max_workers: int = 0,
    progress_callback: Optional[Callable[[int, int], None]] = None,
) -> HitmapResult:
    """
    High-level entrypoint used by the GUI.

    Parameters
    ----------
    data_folder:
        Folder containing *_Mu.txt files.
    config_folder:
        Folder containing detector_config.json, mutxt_config.json,
        analysis_config.json.
    filenames:
        Subset of files (basename only) to process.
    max_workers:
        Number of worker processes (0 -> os.cpu_count()).
    progress_callback:
        Optional callback(progress_done, progress_total).

    Returns
    -------
    HitmapResult
    """
    from ncuhep.muography.classes import (
        PlaneDetector,
        MuTxtFormat,
        AnalysisConfig,
    )

    if not filenames:
        raise ValueError("No input files specified for compute_hitmaps_parallel.")

    data_folder = str(data_folder)
    cfg = Path(config_folder)
    det_path = cfg / "detector_config.json"
    mutxt_path = cfg / "mutxt_config.json"
    ana_path = cfg / "analysis_config.json"

    det = PlaneDetector._import(det_path)
    mutxt = MuTxtFormat._import(mutxt_path)
    ana = AnalysisConfig._import(ana_path)

    layer_count, nx, ny, ncell_layer, ncell_total = build_layer_geometry(det)

    counts_hits_total = np.zeros(ncell_total, dtype=np.int64)
    counts_events_total = np.zeros(ncell_total, dtype=np.int64)
    counts_tracks_total = np.zeros(ncell_total, dtype=np.int64)

    tx_all_list: List[np.ndarray] = []
    ty_all_list: List[np.ndarray] = []
    total_live_time_s = 0.0
    total_events = 0
    total_tracks = 0

    failures: List[Tuple[str, str, float]] = []

    total = len(filenames)
    done = 0

    if max_workers is None or max_workers <= 0:
        max_workers = os.cpu_count() or 1

    logger.info(
        "compute_hitmaps_parallel: data_folder=%s config=%s files=%d workers=%d",
        data_folder,
        config_folder,
        total,
        max_workers,
    )

    with ProcessPoolExecutor(
        max_workers=max_workers,
        initializer=_init_worker,
        initargs=(det, mutxt, ana, data_folder),
    ) as ex:
        futures = {
            ex.submit(_process_one, fn): fn
            for fn in filenames
        }

        for fut in as_completed(futures):
            done += 1
            if progress_callback is not None:
                try:
                    progress_callback(done, total)
                except Exception:
                    logger.exception("progress_callback raised an exception")

            try:
                (
                    fname,
                    ok,
                    err,
                    filesize_mb,
                    h_hits,
                    h_events,
                    h_tracks,
                    tx,
                    ty,
                    live_s,
                    n_ev,
                    n_tr,
                ) = fut.result()
            except Exception as exc:  # noqa: BLE001
                fname = futures[fut]
                failures.append((fname, str(exc), -1.0))
                logger.exception("Worker crashed for file=%s", fname)
                continue

            if not ok:
                failures.append((fname, err, filesize_mb))
                continue

            counts_hits_total += h_hits
            counts_events_total += h_events
            counts_tracks_total += h_tracks

            if tx is not None and ty is not None and tx.size > 0:
                tx_all_list.append(tx)
                ty_all_list.append(ty)

            total_live_time_s += float(live_s)
            total_events += int(n_ev)
            total_tracks += int(n_tr)

    if tx_all_list:
        tx_all = np.concatenate(tx_all_list)
        ty_all = np.concatenate(ty_all_list)
    else:
        tx_all = None
        ty_all = None

    logger.info(
        "compute_hitmaps_parallel summary: Done. Success %d/%d | "
        "Events=%d Tracks=%d | Live time=%.1fs | Failures=%d",
        total - len(failures),
        total,
        total_events,
        total_tracks,
        total_live_time_s,
        len(failures),
    )

    # reshape back to per-layer images
    hits_layers: List[np.ndarray] = []
    events_layers: List[np.ndarray] = []
    tracks_layers: List[np.ndarray] = []

    for layer in range(layer_count):
        start = layer * ncell_layer
        stop = start + ncell_layer

        h = counts_hits_total[start:stop].reshape((ny, nx))
        e = counts_events_total[start:stop].reshape((ny, nx))
        t = counts_tracks_total[start:stop].reshape((ny, nx))

        hits_layers.append(h)
        events_layers.append(e)
        tracks_layers.append(t)

    summary = (
        f"Done. Success {total - len(failures)}/{total} files | "
        f"Events={total_events} Tracks={total_tracks} | "
        f"Live time={total_live_time_s:.1f} s"
    )

    return HitmapResult(
        hits_per_layer=hits_layers,
        events_per_layer=events_layers,
        tracks_per_layer=tracks_layers,
        detector=det,
        summary=summary,
        failures=failures,
        tx_all=tx_all,
        ty_all=ty_all,
        total_live_time_s=total_live_time_s,
        total_events=total_events,
        total_tracks=total_tracks,
    )


# ---------------------------------------------------------------------------
# Flux from in-memory tracks
# ---------------------------------------------------------------------------


def compute_flux_from_tracks(
    tx_all: np.ndarray,
    ty_all: np.ndarray,
    total_live_time_s: float,
    recon_path: str,
    fov_deg: float = 13.0,
):
    """
    Compute flux directly from in-memory track angles (tx_all, ty_all)
    and total live time (seconds), using the provided recon basis.

    This mirrors the standalone script used previously for YMS flux
    calculation, but avoids re-reading *_tracks.npz from disk.
    """
    if tx_all is None or ty_all is None or tx_all.size == 0:
        raise ValueError("No track angles provided for flux computation.")
    if total_live_time_s <= 0.0:
        raise ValueError("Total live time must be > 0 for flux computation.")

    logger.info(
        "compute_flux_from_tracks: Ntracks=%d, time=%.3fs, recon=%s, fov=%g",
        tx_all.size,
        total_live_time_s,
        recon_path,
        fov_deg,
    )

    recon = np.load(recon_path)
    basis = recon["basis"]
    measured_angles = recon["measured_angles"]
    geom_array = recon["geometric_factor"]

    geometric_factor = GeometricFactor()
    geometric_factor.m2_sr = geom_array

    ux = np.unique(measured_angles[0])
    uy = np.unique(measured_angles[1])

    # counts array has the same shape as measured_angles[0]
    counts = np.zeros_like(measured_angles[0], dtype=np.float64)

    # Vectorized binning: tx,ty (rad) -> mrad grid used by recon (round to 0.1 mrad)
    tx_scaled = np.round(tx_all * 1000.0, 1)
    ty_scaled = np.round(ty_all * 1000.0, 1)

    ix = np.searchsorted(ux, tx_scaled)
    iy = np.searchsorted(uy, ty_scaled)

    valid_mask = (
        (ix >= 0) & (ix < len(ux)) &
        (iy >= 0) & (iy < len(uy))
    )
    ix_valid = ix[valid_mask]
    iy_valid = iy[valid_mask]

    np.add.at(counts, (ix_valid, iy_valid), 1)

    total_tracks = tx_all.size
    valid_tracks = ix_valid.size

    summary = (
        f"Flux: total tracks={total_tracks}, used={valid_tracks}, "
        f"time={total_live_time_s:.1f}s"
    )
    logger.info("compute_flux_from_tracks summary: %s", summary)

    # Build N(θx, θy) and U(θx, θy) in basis space
    N = Counts()
    U = Counts()

    N.counts = np.zeros_like(basis[0, 0], dtype=np.float64)

    for i in range(len(ux)):
        for j in range(len(uy)):
            s = np.sum(basis[i, j])
            if s == 0:
                continue
            weight = counts[i, j] / s
            N.counts += basis[i, j] * weight
            U.counts += basis[i, j] * np.sqrt(counts[i, j]) / s

    # Crop to field-of-view
    def crop(image, deg):
        mrad = int(np.radians(deg) * 1000)
        cx = image.shape[1] // 2
        cy = image.shape[0] // 2
        return image[cy - mrad:cy + mrad + 1, cx - mrad:cx + mrad + 1]

    geometric_factor.m2_sr = crop(geometric_factor.m2_sr, fov_deg)
    N.counts = crop(N.counts, fov_deg)
    U.counts = crop(U.counts, fov_deg)

    # Compute flux (unsmoothed)
    timeelapsed = Time()
    timeelapsed.s = float(total_live_time_s)

    flux = N / (geometric_factor * timeelapsed)
    uncertainty = U / (geometric_factor * timeelapsed)

    flux_raw = np.array(flux.counts_m2_s_sr, copy=True)
    unc_raw = np.array(uncertainty.counts_m2_s_sr, copy=True)

    logger.debug(
        "compute_flux_from_tracks: flux shape=%s, uncertainty shape=%s, unit=%s",
        flux_raw.shape,
        unc_raw.shape,
        flux.unit,
    )
    return flux_raw, unc_raw, fov_deg, str(flux.unit), summary
