#!/usr/bin/env python3
"""
UV-Vis Spectrum Simulator for testing the SciLink BO loop.

Usage:
  # 1. Generate initial batch (random or grid)
  python simulate_spectra.py init --output_dir ./spectra --n 8

  # 2. Generate new spectra from BO-suggested parameters
  python simulate_spectra.py run --output_dir ./spectra \
      --params '{"temperature_C": 20.0, "pH": 10.6}'

  # 3. Generate a batch of spectra from BO suggestions
  python simulate_spectra.py run --output_dir ./spectra \
      --params '[{"temperature_C": 20, "pH": 10.6}, {"temperature_C": 35, "pH": 5.2}]'

Physics model (same as generate_test_data.py):
  - Peak center:  450 + 0.3*(T-25) - 5.0*(pH-7)  nm
  - Peak width:   30 + 0.15*(T-25)  nm
  - Peak height:  1.2 - 0.005*(T-25) + 0.04*(pH-7)  (clamped >= 0.2)
  - Shoulder at 400 nm, baseline slope, Gaussian noise

The "ground truth" optimum for peak absorbance is at low temperature + high pH.
"""

import argparse
import json
import sys
import numpy as np
import pandas as pd
from pathlib import Path


# ── Physics model ────────────────────────────────────────────────

WAVELENGTHS = np.linspace(350, 600, 251)

# Valid parameter ranges
BOUNDS = {
    "temperature_C": (5.0, 100.0),
    "pH": (1.0, 14.0),
}


def generate_spectrum(temp: float, ph: float, noise_level: float = 0.005) -> np.ndarray:
    """Generate a synthetic UV-Vis spectrum with physically-motivated shifts."""
    peak_center = 450 + 0.3 * (temp - 25) - 5.0 * (ph - 7.0)
    peak_width = 30 + 0.15 * (temp - 25)
    peak_height = 1.2 - 0.005 * (temp - 25) + 0.04 * (ph - 7.0)
    peak_height = max(peak_height, 0.2)

    absorbance = peak_height * np.exp(
        -0.5 * ((WAVELENGTHS - peak_center) / peak_width) ** 2
    )
    # Shoulder peak
    absorbance += 0.15 * np.exp(
        -0.5 * ((WAVELENGTHS - 400) / 15) ** 2
    )
    # Baseline + noise
    baseline = 0.02 + 0.0001 * (WAVELENGTHS - 350)
    noise = np.random.normal(0, noise_level, len(WAVELENGTHS))

    return np.clip(absorbance + baseline + noise, 0, None)


def save_spectrum(temp: float, ph: float, output_dir: Path, label: str = None):
    """Generate one spectrum, save CSV + sidecar JSON. Returns the filename."""
    absorbance = generate_spectrum(temp, ph)

    if label is None:
        label = f"spectrum_T{temp:.1f}_pH{ph:.1f}"
    # Sanitize for filesystem
    label = label.replace(" ", "_")

    csv_path = output_dir / f"{label}.csv"
    json_path = output_dir / f"{label}.json"

    df = pd.DataFrame({"wavelength_nm": WAVELENGTHS, "absorbance": absorbance})
    df.to_csv(csv_path, index=False)

    conditions = {"temperature_C": round(temp, 2), "pH": round(ph, 2)}
    with open(json_path, "w") as f:
        json.dump(conditions, f, indent=2)

    return csv_path.name


# ── Commands ─────────────────────────────────────────────────────

def cmd_init(args):
    """Generate an initial batch of spectra with random or grid conditions."""
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    n = args.n
    if args.grid:
        # Even grid over the bounds
        n_temp = max(2, int(np.sqrt(n)))
        n_ph = max(2, n // n_temp)
        temps = np.linspace(25, 75, n_temp)
        phs = np.linspace(4, 10, n_ph)
        conditions = [
            {"temperature_C": float(t), "pH": float(p)}
            for t in temps for p in phs
        ][:n]
    else:
        # Random within bounds
        np.random.seed(args.seed)
        conditions = [
            {
                "temperature_C": round(np.random.uniform(20, 80), 1),
                "pH": round(np.random.uniform(3, 11), 1),
            }
            for _ in range(n)
        ]

    global_conditions = {}
    for cond in conditions:
        fname = save_spectrum(cond["temperature_C"], cond["pH"], output_dir)
        global_conditions[fname] = cond
        print(f"  Generated: {fname}")

    # Global conditions file
    with open(output_dir / "conditions.json", "w") as f:
        json.dump(global_conditions, f, indent=2)

    print(f"\nGenerated {len(conditions)} initial spectra in {output_dir}/")
    print(f"Conditions saved to {output_dir}/conditions.json")


def cmd_run(args):
    """Generate spectra for BO-suggested parameters."""
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Parse params
    try:
        params = json.loads(args.params)
    except json.JSONDecodeError as e:
        print(f"Error: could not parse --params JSON: {e}", file=sys.stderr)
        sys.exit(1)

    # Normalize to list
    if isinstance(params, dict):
        params = [params]

    if not isinstance(params, list):
        print("Error: --params must be a JSON dict or list of dicts", file=sys.stderr)
        sys.exit(1)

    # Load existing conditions file to append
    conditions_path = output_dir / "conditions.json"
    if conditions_path.exists():
        with open(conditions_path) as f:
            global_conditions = json.load(f)
    else:
        global_conditions = {}

    # Determine step number from existing files
    existing = list(output_dir.glob("spectrum_*.csv"))
    step = len(existing) + 1

    generated = []
    for i, p in enumerate(params):
        temp = p.get("temperature_C")
        ph = p.get("pH")

        if temp is None or ph is None:
            print(f"Error: entry {i} missing 'temperature_C' or 'pH': {p}", file=sys.stderr)
            continue

        # Clamp to valid ranges
        temp = float(np.clip(temp, *BOUNDS["temperature_C"]))
        ph = float(np.clip(ph, *BOUNDS["pH"]))

        label = f"spectrum_T{temp:.1f}_pH{ph:.1f}"
        fname = save_spectrum(temp, ph, output_dir, label=label)
        global_conditions[fname] = {"temperature_C": round(temp, 2), "pH": round(ph, 2)}
        generated.append(fname)
        print(f"  Generated: {fname}  (T={temp:.1f}C, pH={ph:.1f})")
        step += 1

    # Update global conditions
    with open(conditions_path, "w") as f:
        json.dump(global_conditions, f, indent=2)

    print(f"\nGenerated {len(generated)} new spectra in {output_dir}/")
    print(f"Updated {conditions_path}")

    # Print file paths for easy copy-paste into SciLink
    if generated:
        print(f"\nFiles to analyze:")
        for fname in generated:
            print(f"  {output_dir / fname}")


# ── CLI ──────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="UV-Vis Spectrum Simulator for SciLink BO loop testing"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # init
    p_init = sub.add_parser("init", help="Generate initial batch of spectra")
    p_init.add_argument("--output_dir", required=True, help="Directory for output files")
    p_init.add_argument("--n", type=int, default=8, help="Number of spectra (default: 8)")
    p_init.add_argument("--grid", action="store_true", help="Use grid instead of random")
    p_init.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")

    # run
    p_run = sub.add_parser("run", help="Generate spectra for BO-suggested params")
    p_run.add_argument("--output_dir", required=True, help="Directory for output files")
    p_run.add_argument(
        "--params", required=True,
        help='JSON dict or list: \'{"temperature_C": 20, "pH": 10.6}\''
    )

    args = parser.parse_args()
    if args.command == "init":
        cmd_init(args)
    elif args.command == "run":
        cmd_run(args)


if __name__ == "__main__":
    main()
