#!/usr/bin/env python3
"""
jaxfne-signals-probe-objective-chain — Signals→Probe→Objective analysis stage

Demonstrates the full chain from simulated signals through probe extraction
to objective composition and evaluation.

    signals = model.simulate(...)
           ↓
    probe = signals.probe(...)
           ↓
    objective = jtfne.rate_targets(...) + jtfne.band_power(...)
           ↓
    score = objective(signals)

Every stage is composable and inspectable.
"""

# ============================================================================
# LEVEL 1: SIGNALS OBJECT
# ============================================================================

"""
signals = model.simulate(paradigm=None, seed=0, **sim_kwargs)

Signals is the primary output of simulation. It contains all recorded variables:
  - voltage time series
  - spike rasters
  - extracellular potentials (LFP, CSD, EEG, MEG)
  - network state (currents, synaptic conductances)

Signals attributes (read-only):
  signals.V_m               → array (n_trials, n_steps, n_neurons), voltage in mV
  signals.spikes            → array (n_trials, n_steps, n_neurons), spike counts
  signals.source            → SourceOutput, native current sources
  signals.lfp_contacts      → array (n_trials, n_steps, n_probes), LFP mV
  signals.csd_contacts      → array (n_trials, n_steps, n_probes), CSD μA/mm²
  signals.eeg               → array (n_trials, n_steps, n_channels), scalp EEG mV
  signals.meg               → array (n_trials, n_steps, n_channels), MEG fT

Metadata:
  signals.metadata          → dict {areas, layers, neurons_df, contact_positions, ...}
  signals.shape             → dict {n_trials, n_steps, n_neurons, n_contacts, ...}
"""

# ============================================================================
# LEVEL 1.5: SIGNALS OPERATORS (immediate analysis)
# ============================================================================

"""
Before probing or objective evaluation, analyze signals directly:

1. Firing rates (per neuron, per layer, per cell type):

   rate_all = signals.rate()              # array (n_trials, n_neurons), Hz
   rate_by_layer = signals.rate(by="layer")  # dict {"L1": ..., "L2": ..., ...}
   rate_by_celltype = signals.rate(by="cell_type")  # dict {"E": ..., "PV": ..., ...}

2. Spike statistics:

   cv = signals.cv_isi()                  # coefficient of variation of ISIs
   burst_prop = signals.burst_fraction()  # fraction of spikes in bursts
   irregularity = signals.lv()            # local variation of ISIs

3. Power spectral density:

   psd = signals.psd(fmin=1, fmax=100, resolution=1)  # (freq, power_per_neuron)
   psd_by_layer = signals.psd(by="layer", ...)

4. Bandpower:

   alpha_power = signals.bandpower(band="alpha")      # (n_trials, n_neurons)
   gamma_power = signals.bandpower(band="gamma", fmin=40, fmax=80)

5. Phase (instantaneous):

   phase = signals.phase(frequency_hz=40)  # (n_trials, n_steps, n_neurons)

6. Coherence (cross-frequency):

   coherence = signals.coherence(f1=10, f2=40)  # (n_trials, n_neurons, n_neurons)

7. Summary:

   summary = signals.summary()
   # Output: {
   #   "mean_rate_hz": 8.5,
   #   "rate_by_layer": {"L1": 4.0, "L2": 10.0, ...},
   #   "gamma_power": 0.15,
   #   "alpha_beta_power": 0.25,
   #   "burstiness": 0.3,
   #   "synchrony": 0.45,
   # }

8. Validation (check for NaN, Inf, physiological bounds):

   assert signals.validate()  # raises if issues found
   issues = signals.validate(verbose=True)  # returns list of issues
"""

# ============================================================================
# LEVEL 2: PROBE EXTRACTION (Emitter → Source → Field → Probe)
# ============================================================================

"""
TFNE doctrine says:
  Emitter (neurons firing)
    → Source (transmembrane current)
      → Field (extracellular potential)
        → Probe (measured signal)

signals.probe(kind, **kwargs) extracts different probe types:

1. LFP-proxy (local field potential, Gaussian kernel + CSD):

   lfp_probe = signals.probe("LFP-proxy")
   # Returns: LFPProbe object
   # Data: lfp_probe.signal (n_trials, n_steps, n_contacts)
   # Metadata: lfp_probe.contact_positions, lfp_probe.contact_labels

2. CSD-proxy (current source density, spatial second derivative of LFP):

   csd_probe = signals.probe("CSD-proxy")
   # Automatically computed from LFP
   # High CSD = local current sink/source

3. EEG-proxy (scalp electrodes, dipole model + lead-field):

   eeg_probe = signals.probe("EEG-proxy")
   # Simulated 10–20 electrode placement
   # Much lower amplitude than LFP (far-field)

4. MEG-proxy (magnetoencephalography, oriented dipoles + lead-field):

   meg_probe = signals.probe("MEG-proxy")
   # Sensitive to tangential dipoles only
   # Oriented lead-field matrix

5. MUA-proxy (multi-unit activity, spike detection from voltage):

   mua_probe = signals.probe("MUA-proxy")
   # Extracts spike times in defined windows
   # Returns event times, not continuous signal

6. EMM (extracellular monopole moment, native transmembrane current):

   emm_probe = signals.probe("EMM")
   # Direct measure of network current flow
   # No filtering, full resolution

Probe methods:

   probe.summary()           → dict (probe_type, n_contacts, frequency_content, ...)
   probe.spectrum(fmin=1, fmax=100)  → (freq, power_per_contact)
   probe.phase(frequency_hz)  → (n_trials, n_steps, n_contacts)
   probe.coherence()         → cross-contact coherence matrix
   probe.plot(freq_range=[1, 100], cmap="viridis")  → Figure
   probe.export(path, format="HDF5")  → save to disk

Example:

   lfp_probe = signals.probe("LFP-proxy")
   lfp_spectrum = lfp_probe.spectrum(fmin=1, fmax=100)
   lfp_coherence = lfp_probe.coherence()
   fig = lfp_probe.plot(freq_range=[1, 100])
   fig.savefig("lfp_spectrogram.png")
"""

# ============================================================================
# LEVEL 2.5: SPECTROLAMINAR ANALYSIS
# ============================================================================

"""
Special case: 1D laminar probes (vertically stacked contacts).
Spectrolaminar = frequency × depth decomposition.

from jaxfne.vis.tutorial_panels import spectrolaminar_suite_3panel

probe = signals.probe("LFP-proxy")
scores, specs = tu.summarize_spectrolaminar_similarity(signals, cfg)

# specs dict contains per-layer, per-frequency relative power
# canonical spectrolaminar pattern:
#   - Alpha/beta (10–30 Hz) strong in L2/3 (superficial)
#   - Gamma (40–100 Hz) strong in L5/6 (deep)
#   - Theta/delta (1–10 Hz) widespread

figs = spectrolaminar_suite_3panel(
    specs, model, cfg,
    areas=["V1"],
    stage="cortical_column_spectrolaminar",
    output_dir="./figs",
    theme="dark",
    profile_smooth_sigma=1.0,
)

# Returns: {area: Figure} with 3-panel depth×freq heatmap + power profiles + CSD
"""

# ============================================================================
# LEVEL 3: OBJECTIVE CONSTRUCTORS (direct definitions)
# ============================================================================

"""
Build objectives to score simulation outcomes.

1. Rate targets (firing rate per layer, per cell type):

   obj_rate = jtfne.rate_targets(
       groups={"L2": 10, "L5": 8, "L6": 5},  # target Hz
       weights={"L2": 1.0, "L5": 0.8, "L6": 0.5},  # relative importance
       loss_fn="mse",  # "mse" or "l1"
   )

2. Band power targets (spectral energy in frequency band):

   obj_gamma = jtfne.band_power(
       band="gamma",           # or: fmin=40, fmax=80
       target=0.25,            # target relative power
       layers=["L2/3"],        # restrict to layer(s)
       weight=1.0,
   )

3. Phase locking (synchrony within frequency band):

   obj_phase = jtfne.phase_locking(
       frequency_hz=40,        # or: band="gamma"
       target_kappa=0.5,       # target concentration (0=uniform, 1=perfect sync)
       layers=["L2/3"],
       weight=1.0,
   )

4. CSD pattern matching (match a canonical CSD profile):

   target_csd = np.array([...])  # canonical CSD pattern
   obj_csd = jtfne.csd_pattern(
       target_pattern=target_csd,
       weight=1.0,
   )

5. Evoked response (match a target response to stimulus):

   obj_evoked = jtfne.evoked_response(
       target_waveform=target_vm,  # target voltage response
       time_window=[0, 100],  # ms post-stimulus
       weight=1.0,
   )
"""

# ============================================================================
# LEVEL 3.5: OBJECTIVE ALGEBRA (composition)
# ============================================================================

"""
Objectives can be combined using algebra:

obj1 = jtfne.rate_targets({"L2": 10})
obj2 = jtfne.band_power(band="gamma", target=0.25)
obj3 = jtfne.phase_locking(frequency_hz=40, target_kappa=0.5)

# Sum (weighted combination):
objective = obj1 + obj2 + obj3
# Loss = loss1(signals) + loss2(signals) + loss3(signals)

# With custom weights:
objective = obj1 + obj2 * 0.5 + obj3 * 0.3
# Loss = loss1 + 0.5*loss2 + 0.3*loss3

# Negation (maximize instead of minimize):
objective = obj1 - obj2  # maximize obj1, minimize obj2

# Scalar operations:
objective = (obj1 + obj2) * 0.5  # rescale entire composition

Example (realistic Phase 0 objective):

objective = (
    jtfne.rate_targets(
        {"L1": 5, "L2": 10, "L3": 8, "L4": 7, "L5": 8, "L6": 5},
        weights={"L2": 1.2, "L3": 1.0, "L5": 0.8},  # prefer L2/3 activity
    )
    + jtfne.band_power(band="gamma", target=0.25, layers=["L2/3"]) * 0.5
    + jtfne.phase_locking(frequency_hz=40, target_kappa=0.4) * 0.3
)
"""

# ============================================================================
# LEVEL 4: OBJECTIVE EVALUATION
# ============================================================================

"""
Evaluate objectives on signals.

1. Quick evaluation (scalar score):

   score = objective(signals)
   # Returns: float (total loss)
   # Used during optimization

2. Detailed evaluation (breakdown by component):

   report = objective.evaluate(signals)
   # Returns: ObjectiveReport
   # Contains:
   #   report.total_loss     → float
   #   report.losses         → dict {obj_name: loss_value, ...}
   #   report.details        → dict {obj_name: detailed_breakdown, ...}
   #   report.summary()      → str (human-readable report)

3. Per-component inspection:

   for obj_name, loss_val in report.losses.items():
       print(f"{obj_name}: {loss_val:.4f}")

Example:

signals = model.simulate(seed=0)

objective = (
    jtfne.rate_targets({"L2": 10, "L5": 8})
    + jtfne.band_power(band="gamma", target=0.25) * 0.5
)

score = objective(signals)
print(f"Total loss: {score:.4f}")  # single number

report = objective.evaluate(signals)
print(report.summary())
# Output:
# ============================================
# Objective Report
# ============================================
# Total loss: 1.2345
#
# Component losses:
#   rate_targets:        0.8234
#   band_power (γ):      0.4111
#
# Details:
#   rate_targets:
#     L2 rate:         10.1 Hz (target: 10.0)
#     L5 rate:          7.8 Hz (target: 8.0)
#   band_power (γ):
#     γ power:         0.24 (target: 0.25)
# ============================================
"""

# ============================================================================
# COMPLETE ANALYSIS CHAIN EXAMPLE
# ============================================================================

"""
from pathlib import Path
import jaxfne as jtfne
from jaxfne import tutorial_utils as tu
from jaxfne.vis.tutorial_panels import spectrolaminar_suite_3panel

# ─────────────────────────────────────────────────────────────────────────
# Setup & Config (see jaxfne-configuration-fluent-api skill)
# ─────────────────────────────────────────────────────────────────────────

LAYER_FRACTIONS = {
    "L1": (0.0, 0.1),    "L2": (0.1, 0.35),   "L3": (0.35, 0.55),
    "L4": (0.55, 0.65),  "L5": (0.65, 0.85), "L6": (0.85, 1.0),
}

LAYER_CELL_TYPE_FRAC = {
    "L1": {"E": 0.35, "PV": 0.30, "SST": 0.20, "VIP": 0.15},
    "L2": {"E": 0.45, "PV": 0.25, "SST": 0.18, "VIP": 0.12},
    "L3": {"E": 0.55, "PV": 0.20, "SST": 0.15, "VIP": 0.10},
    "L4": {"E": 0.65, "PV": 0.18, "SST": 0.12, "VIP": 0.05},
    "L5": {"E": 0.85, "PV": 0.08, "SST": 0.05, "VIP": 0.02},
    "L6": {"E": 0.90, "PV": 0.06, "SST": 0.03, "VIP": 0.01},
}

layer_fracs = tuple(
    (layer, LAYER_FRACTIONS[layer][0], LAYER_FRACTIONS[layer][1])
    for layer in ("L1", "L2", "L3", "L4", "L5", "L6")
)

cfg = tu.make_laminar_column_config(
    areas=("V1",),
    layers=("L1", "L2", "L3", "L4", "L5", "L6"),
    layer_fractions=layer_fracs,
    layer_cell_type_frac=LAYER_CELL_TYPE_FRAC,
    n_neuron_per_column=1000,
    duration_ms=1000.0, dt_ms=0.1,
    n_trials=4, n_contacts=64, freq_count=96,
    seed=0,
)

model = tu.build_laminar_column(cfg)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 1: Simulate → Signals
# ─────────────────────────────────────────────────────────────────────────

signals = model.simulate(seed=0)
print(f"Simulation complete: {signals.shape}")

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 1.5: Immediate signal analysis
# ─────────────────────────────────────────────────────────────────────────

rate_by_layer = signals.rate(by="layer")
print(f"Firing rates by layer: {rate_by_layer}")

gamma_power = signals.bandpower(band="gamma")
print(f"Gamma power: {gamma_power.mean():.3f}")

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 2: Probe extraction
# ─────────────────────────────────────────────────────────────────────────

lfp_probe = signals.probe("LFP-proxy")
csd_probe = signals.probe("CSD-proxy")
eeg_probe = signals.probe("EEG-proxy")

print(f"LFP probe: {lfp_probe.shape}")
print(f"EEG probe: {eeg_probe.shape}")

# Spectrolaminar (canonical profile):
scores, specs = tu.summarize_spectrolaminar_similarity(signals, cfg)
figs = spectrolaminar_suite_3panel(
    specs, model, cfg,
    areas=["V1"],
    stage="canonical_cortical_column",
    output_dir="./figs",
    theme="dark",
)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 3: Objectives (what we want to optimize)
# ─────────────────────────────────────────────────────────────────────────

objective = (
    jtfne.rate_targets(
        {"L2": 10, "L3": 8, "L5": 8, "L6": 5},
        weights={"L2": 1.2, "L3": 1.0, "L5": 0.8, "L6": 0.5},
    )
    + jtfne.band_power(band="gamma", target=0.25, layers=["L2/3"]) * 0.5
    + jtfne.phase_locking(frequency_hz=40, target_kappa=0.4) * 0.3
)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 4: Evaluate objectives on current signals
# ─────────────────────────────────────────────────────────────────────────

score = objective(signals)
print(f"Current objective score: {score:.4f}")

report = objective.evaluate(signals)
print(report.summary())

# ─────────────────────────────────────────────────────────────────────────
# Next: Pass to optimizer (see jaxfne-optimizer-grammar skill)
# ─────────────────────────────────────────────────────────────────────────

# optimizer = jtfne.AGSDR(learning_rate=0.01)
# result = optimizer.optimize(model, objective, n_steps=500, ...)
"""

# ============================================================================
# RULE
# ============================================================================

"""
[INVARIANT] Every Signals analysis must flow through this chain:

  ✅ CORRECT:
    signals = model.simulate(...)
    probe = signals.probe("LFP-proxy")
    objective = jtfne.rate_targets(...) + jtfne.band_power(...)
    score = objective(signals)

  ❌ WRONG:
    signals = model.simulate(...)
    lfp = signals.lfp_contacts  # direct attribute access
    manual_psd = np.fft.rfft(lfp[0, :, 0])  # hand-rolled FFT
    manual_rate = np.sum(signals.spikes) / n_steps  # direct indexing
    # (bypasses validation, composability, TFNE governance)

Why?
  - Probes encode the Emitter→Source→Field→Probe doctrine
  - Objectives are composable and algebraic
  - Evaluation reports provide reproducible audits
  - Skipping stages loses TFNE metadata and traceability
"""
