"""
VMS Slit Experiment – Image Analysis Helper

This script (prototype) was used to generate the per-frame metrics
reported in results_set1.rtf for the images in the ZIP archives
(centroid, FWHM radius, inner/outer flux, etc.).

It is provided for transparency and replication; it is not a polished
or fully automated pipeline. Some images required manual steps or
separate passes that are not included here.
"""

#!/usr/bin/env python3
"""
VMS Slit Experiment – Image Analysis Helper

Prototype script used to generate per-frame metrics for the slit experiment.

- Reads a manifest CSV with columns:
    file_name,block_id,distance_nominal_mm,frame_role
- Loads each image from a given folder.
- Measures:
    centroid_x_px, centroid_y_px,
    peak_intensity,
    fwhm_radius_px,
    w10_radius_px,
    inner_flux, outer_flux, outer_fraction,
    lobe_count (very approximate)

- For each (block_id, distance_nominal_mm) group, one frame whose frame_role
  contains the word "baseline" is treated as the baseline.
  Deltas are computed relative to that baseline:

    inner_flux_delta     = inner_flux - baseline_inner_flux
    outer_fraction_delta = outer_fraction - baseline_outer_fraction
    fwhm_delta_px        = fwhm_radius_px - baseline_fwhm_radius_px

This is a **prototype** and not a polished pipeline. Thresholds, masks and the
radial-profile logic are intentionally simple and may need tuning for other
cameras or configurations.

Dependencies:
    pip install numpy opencv-python scipy pandas

Usage:
    python VMS_Slit_Image_Analysis.py \
        --image-dir PATH/TO/IMAGES \
        --manifest  PATH/TO/manifest.csv \
        --output    results.txt
"""

import argparse
import math
import os
from dataclasses import dataclass, asdict
from typing import Dict, Tuple, List

import cv2
import numpy as np
import pandas as pd
from scipy.ndimage import gaussian_filter1d

@dataclass
class FrameInfo:
    file_name: str
    block_id: int
    distance_nominal_mm: float
    frame_role: str

@dataclass
class FrameResult:
    file_name: str
    block_id: int
    distance_nominal_mm: float
    frame_role: str

    centroid_x_px: float
    centroid_y_px: float
    drift_dx_from_block: float
    drift_dy_from_block: float

    peak_intensity: float
    fwhm_radius_px: float
    w10_radius_px: float

    inner_flux: float
    outer_flux: float
    outer_fraction: float
    lobe_count: int

    inner_flux_delta: float
    outer_fraction_delta: float
    fwhm_delta_px: float

def load_manifest(path: str) -> List[FrameInfo]:
    df = pd.read_csv(path)
    required = {"file_name", "block_id", "distance_nominal_mm", "frame_role"}
    missing = required.difference(df.columns)
    if missing:
        raise ValueError(f"Manifest is missing columns: {missing}")

    infos: List[FrameInfo] = []
    for _, row in df.iterrows():
        infos.append(
            FrameInfo(
                file_name=str(row["file_name"]),
                block_id=int(row["block_id"]),
                distance_nominal_mm=float(row["distance_nominal_mm"]),
                frame_role=str(row["frame_role"]),
            )
        )
    return infos

def load_image_gray(image_path: str) -> np.ndarray:
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(f"Could not read image: {image_path}")
    return img.astype(np.float32)

def compute_centroid_and_peak(img: np.ndarray) -> Tuple[float, float, float]:
    """Intensity-weighted centroid and peak intensity."""
    h, w = img.shape
    y_indices, x_indices = np.indices((h, w))

    total_intensity = img.sum()
    if total_intensity <= 0:
        # Degenerate, avoid divide by zero
        cx = w / 2.0
        cy = h / 2.0
    else:
        cx = (img * x_indices).sum() / total_intensity
        cy = (img * y_indices).sum() / total_intensity

    peak = float(img.max())
    return float(cx), float(cy), peak

def radial_profile(img: np.ndarray, cx: float, cy: float) -> Tuple[np.ndarray, np.ndarray]:
    """Compute azimuthally averaged radial intensity profile."""
    h, w = img.shape
    y_indices, x_indices = np.indices((h, w))
    r = np.sqrt((x_indices - cx) ** 2 + (y_indices - cy) ** 2)
    r_flat = r.flatten()
    i_flat = img.flatten()

    # Bin radii to integer pixels
    r_int = r_flat.astype(int)
    max_r = r_int.max()
    radial_sum = np.bincount(r_int, weights=i_flat, minlength=max_r + 1)
    radial_count = np.bincount(r_int, minlength=max_r + 1)
    radial_count[radial_count == 0] = 1  # avoid /0
    profile = radial_sum / radial_count
    radii = np.arange(len(profile), dtype=float)
    return radii, profile

def find_radius_at_fraction(
    radii: np.ndarray, profile: np.ndarray, fraction: float
) -> float:
    """Find radius where profile falls to given fraction of its max."""
    if profile.size == 0:
        return 0.0
    peak = profile.max()
    target = peak * fraction
    # find first index where profile <= target after the maximum
    peak_idx = int(np.argmax(profile))
    segment_r = radii[peak_idx:]
    segment_p = profile[peak_idx:]
    below = np.where(segment_p <= target)[0]
    if len(below) == 0:
        return float(segment_r[-1])
    idx = below[0]
    if idx == 0:
        return float(segment_r[0])
    # linear interpolation between segment_r[idx-1] and segment_r[idx]
    r1, r2 = segment_r[idx - 1], segment_r[idx]
    p1, p2 = segment_p[idx - 1], segment_p[idx]
    if p2 == p1:
        return float(r2)
    t = (target - p1) / (p2 - p1)
    return float(r1 + t * (r2 - r1))

def count_lobes(profile: np.ndarray, min_fraction: float = 0.1) -> int:
    """
    Crude lobe counter: smooth profile, then count local maxima above
    min_fraction * max(profile).
    """
    if profile.size == 0:
        return 0
    smoothed = gaussian_filter1d(profile, sigma=2.0)
    peak = smoothed.max()
    threshold = peak * min_fraction
    lobes = 0
    for i in range(1, len(smoothed) - 1):
        if smoothed[i] > smoothed[i - 1] and smoothed[i] > smoothed[i + 1]:
            if smoothed[i] >= threshold:
                lobes += 1
    return int(lobes)

def compute_fluxes(img: np.ndarray, cx: float, cy: float, r_fwhm: float, r_w10: float):
    """Sum intensity inside r_fwhm and between r_fwhm..r_w10."""
    h, w = img.shape
    y_indices, x_indices = np.indices((h, w))
    r = np.sqrt((x_indices - cx) ** 2 + (y_indices - cy) ** 2)

    inner_mask = r <= r_fwhm
    outer_mask = (r > r_fwhm) & (r <= r_w10)

    inner_flux = float(img[inner_mask].sum())
    outer_flux = float(img[outer_mask].sum())
    total = inner_flux + outer_flux
    if total <= 0:
        outer_fraction = 0.0
    else:
        outer_fraction = outer_flux / total

    return inner_flux, outer_flux, outer_fraction

def analyze_frame(
    img: np.ndarray,
    info: FrameInfo,
    block_center: Tuple[float, float],
    baseline_metrics: Dict[str, float],
) -> FrameResult:
    cx, cy, peak = compute_centroid_and_peak(img)
    drift_dx = cx - block_center[0]
    drift_dy = cy - block_center[1]

    radii, profile = radial_profile(img, cx, cy)
    r_fwhm = find_radius_at_fraction(radii, profile, 0.5)
    r_w10 = find_radius_at_fraction(radii, profile, 0.1)
    lobes = count_lobes(profile)

    inner_flux, outer_flux, outer_fraction = compute_fluxes(
        img, cx, cy, r_fwhm, r_w10
    )

    # Deltas vs baseline for this (block_id, distance) group
    inner_flux_delta = inner_flux - baseline_metrics["inner_flux"]
    outer_fraction_delta = outer_fraction - baseline_metrics["outer_fraction"]
    fwhm_delta_px = r_fwhm - baseline_metrics["fwhm_radius_px"]

    return FrameResult(
        file_name=info.file_name,
        block_id=info.block_id,
        distance_nominal_mm=info.distance_nominal_mm,
        frame_role=info.frame_role,
        centroid_x_px=cx,
        centroid_y_px=cy,
        drift_dx_from_block=drift_dx,
        drift_dy_from_block=drift_dy,
        peak_intensity=peak,
        fwhm_radius_px=r_fwhm,
        w10_radius_px=r_w10,
        inner_flux=inner_flux,
        outer_flux=outer_flux,
        outer_fraction=outer_fraction,
        lobe_count=lobes,
        inner_flux_delta=inner_flux_delta,
        outer_fraction_delta=outer_fraction_delta,
        fwhm_delta_px=fwhm_delta_px,
    )

def print_result_block(result: FrameResult, fh):
    """
    Print a single result block in the format:

    file_name            = IMG_1268.jpeg
    block_id             = 7
    ...
    ------------------------------------------------------------
    """
    d = asdict(result)
    fh.write(f"file_name            = {d['file_name']}\n")
    fh.write(f"block_id             = {d['block_id']}\n")
    fh.write(f"distance_nominal_mm  = {d['distance_nominal_mm']:.2f}\n")
    fh.write(f"frame_role           = {d['frame_role']}\n\n")

    fh.write(f"centroid_x_px        = {d['centroid_x_px']:.2f}\n")
    fh.write(f"centroid_y_px        = {d['centroid_y_px']:.2f}\n")
    fh.write(f"drift_dx_from_block  = {d['drift_dx_from_block']:.3f}\n")
    fh.write(f"drift_dy_from_block  = {d['drift_dy_from_block']:.3f}\n\n")

    fh.write(f"peak_intensity       = {d['peak_intensity']:.3f}\n")
    fh.write(f"fwhm_radius_px       = {d['fwhm_radius_px']:.2f}\n")
    fh.write(f"w10_radius_px        = {d['w10_radius_px']:.2f}\n")
    fh.write(f"inner_flux           = {d['inner_flux']:.1f}\n")
    fh.write(f"outer_flux           = {d['outer_flux']:.1f}\n")
    fh.write(f"outer_fraction       = {d['outer_fraction']:.6f}\n")
    fh.write(f"lobe_count           = {d['lobe_count']}\n\n")

    fh.write(f"inner_flux_delta     = {d['inner_flux_delta']:.1f}\n")
    fh.write(
        f"outer_fraction_delta = {d['outer_fraction_delta']:.6f}\n"
    )
    fh.write(f"fwhm_delta_px        = {d['fwhm_delta_px']:.3f}\n\n")
    fh.write("------------------------------------------------------------\n\n")

def main():
    parser = argparse.ArgumentParser(
        description="Prototype slit-image analysis for VMS experiment."
    )
    parser.add_argument(
        "--image-dir", required=True, help="Directory containing JPEG images."
    )
    parser.add_argument(
        "--manifest", required=True, help="CSV manifest describing frames."
    )
    parser.add_argument(
        "--output",
        required=True,
        help="Output text file (results in human-readable blocks).",
    )
    parser.add_argument(
        "--block-center-x",
        type=float,
        default=None,
        help="Optional fixed block center X (px). If omitted, uses frame centroid.",
    )
    parser.add_argument(
        "--block-center-y",
        type=float,
        default=None,
        help="Optional fixed block center Y (px). If omitted, uses frame centroid.",
    )

    args = parser.parse_args()

    infos = load_manifest(args.manifest)

    # Group by (block_id, distance_nominal_mm)
    groups: Dict[Tuple[int, float], List[FrameInfo]] = {}
    for info in infos:
        key = (info.block_id, info.distance_nominal_mm)
        groups.setdefault(key, []).append(info)

    with open(args.output, "w") as fh:
        for (block_id, distance), frames in sorted(groups.items()):
            # Identify baseline frame
            baseline_frame = None
            for info in frames:
                if "baseline" in info.frame_role.lower():
                    baseline_frame = info
                    break
            if baseline_frame is None:
                # Fallback: first frame in group as baseline
                baseline_frame = frames[0]

            baseline_path = os.path.join(args.image_dir, baseline_frame.file_name)
            baseline_img = load_image_gray(baseline_path)

            # block center: either user-specified or baseline centroid
            if args.block_center_x is not None and args.block_center_y is not None:
                block_cx, block_cy = args.block_center_x, args.block_center_y
            else:
                block_cx, block_cy, _ = compute_centroid_and_peak(baseline_img)

            # Baseline metrics
            radii_b, profile_b = radial_profile(baseline_img, block_cx, block_cy)
            r_fwhm_b = find_radius_at_fraction(radii_b, profile_b, 0.5)
            inner_b, outer_b, outer_frac_b = compute_fluxes(
                baseline_img, block_cx, block_cy, r_fwhm_b, find_radius_at_fraction(radii_b, profile_b, 0.1)
            )
            baseline_metrics = {
                "inner_flux": inner_b,
                "outer_fraction": outer_frac_b,
                "fwhm_radius_px": r_fwhm_b,
            }

            for info in frames:
                img_path = os.path.join(args.image_dir, info.file_name)
                img = load_image_gray(img_path)
                result = analyze_frame(
                    img, info, (block_cx, block_cy), baseline_metrics
                )
                print_result_block(result, fh)

if __name__ == "__main__":
    main()