"""
Custom peak detection agent — example of a minimal SciLink custom agent.

Demonstrates:
- Only analyze() needs to be implemented; all other BaseAnalysisAgent
  methods have working defaults.
- No LLM call is required for the core analysis logic.
- The standard result dict format expected by the orchestrator.

Usage with scilink analyze CLI
-------------------------------
  scilink analyze --agents examples/custom_peak_agent.py

Then in the chat session:
  > Analyze examples/data/xps_ti2p.csv using agent 4

Usage programmatically
-----------------------
  from examples.custom_peak_agent import PeakDetectionAgent

  agent = PeakDetectionAgent(output_dir="./output")
  result = agent.analyze("examples/data/xps_ti2p.csv")
  print(result["detailed_analysis"])
"""

import json
import logging
import numpy as np
from pathlib import Path
from typing import Any, Dict, List, Optional

from scilink.agents.exp_agents.base_agent import BaseAnalysisAgent


class PeakDetectionAgent(BaseAnalysisAgent):
    """
    Detects and characterises peaks in 1-D spectral / curve data.

    Supports CSV (two-column: x, y) and NPY (1-D or 2-D [N,2]) files.
    Peak detection uses scipy.signal.find_peaks when scipy is available
    and falls back to a simple numpy-based local-maximum search otherwise.

    This agent intentionally keeps all analysis in Python (no LLM call)
    to serve as a clear, runnable example of the custom-agent API.
    """

    # ── Metadata used by register_agent() and the /agents CLI command ──────
    AGENT_NAME = "PeakDetectionAgent"
    AGENT_DESCRIPTION = (
        "Peak detection and characterisation for 1-D spectra and curves "
        "(XPS, Raman, XRD, UV-Vis, …). Accepts CSV or NPY files."
    )
    AGENT_SHORT_NAME = "PeakDet"

    # ── Tunable defaults ────────────────────────────────────────────────────
    MIN_PEAK_HEIGHT_FRACTION = 0.05   # ignore peaks below 5 % of max intensity
    MIN_PEAK_PROMINENCE_FRACTION = 0.03

    def analyze(
        self,
        data,
        system_info: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Detect peaks in a 1-D dataset and return a standardised result dict.

        Args:
            data:        Path to a CSV or NPY file (str or Path), or a list
                         of such paths (only the first file is processed).
            system_info: Optional metadata dict (used for axis-label hints).
            **kwargs:    Ignored — present for API compatibility.

        Returns:
            dict with keys: status, detailed_analysis, scientific_claims,
            output_directory, peaks (list of dicts).
        """
        # ── Resolve data path ───────────────────────────────────────────────
        if isinstance(data, list):
            data = data[0]
        data_path = Path(data)

        if not data_path.exists():
            return self._error(f"File not found: {data_path}")

        # ── Load data ───────────────────────────────────────────────────────
        try:
            x, y, x_label, y_label = self._load(data_path, system_info or {})
        except Exception as exc:
            return self._error(f"Could not load data: {exc}")

        # ── Detect peaks ────────────────────────────────────────────────────
        peak_indices, properties = self._find_peaks(y)

        if len(peak_indices) == 0:
            summary = "No peaks detected above the prominence threshold."
            peaks_data = []
        else:
            peaks_data = self._describe_peaks(x, y, peak_indices, properties)
            summary = self._build_summary(peaks_data, x_label, y_label)

        # ── Save a simple text report ────────────────────────────────────────
        report_path = self.output_dir / "peak_report.txt"
        report_path.write_text(summary)

        # ── Build scientific claims ─────────────────────────────────────────
        claims = self._build_claims(peaks_data, x_label, y_label, data_path.name)

        return {
            "status": "success",
            "detailed_analysis": summary,
            "scientific_claims": claims,
            "output_directory": str(self.output_dir),
            "peaks": peaks_data,          # extra key; orchestrator ignores it
            "report_path": str(report_path),
        }

    # ── Private helpers ─────────────────────────────────────────────────────

    def _load(self, path: Path, system_info: dict):
        """Load x/y arrays from CSV or NPY and infer axis labels."""
        x_label = system_info.get("x_axis", "x")
        y_label = system_info.get("y_axis", "Intensity")

        if path.suffix.lower() in {".csv", ".txt", ".tsv"}:
            import csv
            rows = []
            header = None
            with open(path, newline="") as fh:
                reader = csv.reader(fh)
                for row in reader:
                    try:
                        rows.append([float(v) for v in row])
                    except ValueError:
                        # treat as header
                        if header is None:
                            header = row
            arr = np.array(rows)
            if header and len(header) >= 2:
                x_label = header[0].strip() or x_label
                y_label = header[1].strip() or y_label
            x, y = arr[:, 0], arr[:, 1]

        elif path.suffix.lower() == ".npy":
            arr = np.load(str(path))
            if arr.ndim == 1:
                x = np.arange(len(arr))
                y = arr
            elif arr.ndim == 2 and arr.shape[1] == 2:
                x, y = arr[:, 0], arr[:, 1]
            elif arr.ndim == 2 and arr.shape[0] == 2:
                x, y = arr[0], arr[1]
            else:
                raise ValueError(f"Unsupported NPY shape: {arr.shape}")
        else:
            raise ValueError(f"Unsupported file type: {path.suffix}")

        return x, y, x_label, y_label

    def _find_peaks(self, y: np.ndarray):
        """Return (peak_indices, properties) using scipy or numpy fallback."""
        height_threshold = float(np.max(y)) * self.MIN_PEAK_HEIGHT_FRACTION
        prominence_threshold = float(np.max(y) - np.min(y)) * self.MIN_PEAK_PROMINENCE_FRACTION

        try:
            from scipy.signal import find_peaks
            indices, props = find_peaks(
                y,
                height=height_threshold,
                prominence=prominence_threshold,
            )
            return indices, props
        except ImportError:
            logging.info("scipy not available — using numpy local-max fallback")
            return self._numpy_peaks(y, height_threshold), {}

    @staticmethod
    def _numpy_peaks(y: np.ndarray, height_threshold: float):
        """Simple local-maximum search (no scipy required)."""
        indices = []
        for i in range(1, len(y) - 1):
            if y[i] > y[i - 1] and y[i] > y[i + 1] and y[i] >= height_threshold:
                indices.append(i)
        return np.array(indices, dtype=int)

    @staticmethod
    def _describe_peaks(x, y, indices, properties) -> List[dict]:
        peaks = []
        prominences = properties.get("prominences", np.full(len(indices), float("nan")))
        for rank, idx in enumerate(indices):
            peaks.append({
                "rank": rank + 1,
                "position": float(x[idx]),
                "intensity": float(y[idx]),
                "prominence": float(prominences[rank]) if rank < len(prominences) else None,
            })
        # Sort by intensity descending
        peaks.sort(key=lambda p: p["intensity"], reverse=True)
        for rank, p in enumerate(peaks):
            p["rank"] = rank + 1
        return peaks

    def _build_summary(self, peaks: List[dict], x_label: str, y_label: str) -> str:
        lines = [
            f"Peak Detection Report",
            f"{'=' * 40}",
            f"Total peaks found: {len(peaks)}",
            "",
            f"{'Rank':<6} {'Position (' + x_label + ')':<28} {'Intensity (' + y_label + ')':<26} {'Prominence':<12}",
            "-" * 75,
        ]
        for p in peaks:
            prom = f"{p['prominence']:.1f}" if p["prominence"] is not None else "n/a"
            lines.append(
                f"{p['rank']:<6} {p['position']:<28.4f} {p['intensity']:<26.1f} {prom:<12}"
            )
        lines += [
            "",
            f"Dominant peak: {peaks[0]['position']:.4f} {x_label} "
            f"(intensity {peaks[0]['intensity']:.1f} {y_label})",
        ]
        return "\n".join(lines)

    @staticmethod
    def _build_claims(peaks, x_label, y_label, filename) -> List[dict]:
        if not peaks:
            return []
        dominant = peaks[0]
        claims = [
            {
                "claim": (
                    f"The spectrum '{filename}' contains {len(peaks)} resolvable peak(s). "
                    f"The dominant peak is located at "
                    f"{dominant['position']:.4f} {x_label} with an intensity of "
                    f"{dominant['intensity']:.1f} {y_label}."
                ),
                "scientific_impact": "Identifies the primary spectral feature for further analysis.",
                "has_anyone_question": (
                    f"Has anyone reported peaks near {dominant['position']:.1f} {x_label} "
                    f"in this type of spectrum?"
                ),
                "keywords": ["peak detection", x_label, y_label, filename],
            }
        ]
        if len(peaks) >= 2:
            second = peaks[1]
            claims.append({
                "claim": (
                    f"A secondary peak is observed at {second['position']:.4f} {x_label} "
                    f"(intensity {second['intensity']:.1f} {y_label}), "
                    f"suggesting the presence of multiple spectral components."
                ),
                "scientific_impact": "Multiple peaks may indicate mixed phases or chemical states.",
                "has_anyone_question": (
                    f"Have multiple peaks near {dominant['position']:.1f} and "
                    f"{second['position']:.1f} {x_label} been reported for similar samples?"
                ),
                "keywords": ["multiple peaks", "chemical states", x_label, filename],
            })
        return claims

    @staticmethod
    def _error(message: str) -> dict:
        return {
            "status": "error",
            "error": {"error": message},
            "detailed_analysis": "",
            "scientific_claims": [],
        }
