#!/usr/bin/env python3
"""
jaxfne-objective-grammar — THE MANDATORY LAW

Every jaxfne code must follow the fluent object-transform grammar:

    A
    A.B(...)
    A.B(...).C(...)
    A.B(...).C(...).D(...)
    ...

where each stage returns a new typed object and remains inspectable, exportable, and composable.

The canonical TFNE chain is:

    Configuration → Model → Signals → Probe → Objective → Optimizer → Manifest

where every arrow represents an object that exposes additional methods.
"""

# ============================================================================
# LEVEL 0: ROOT NAMESPACE
# ============================================================================

"""
import jaxfne as jtfne

Root entry points:
  jtfne.Configuration()      → Configuration object
  jtfne.construct(cfg)       → Model object
  jtfne.simulate(model, ...) → Signals object
  jtfne.optimize(obj, ...)   → Optimizer object
  jtfne.export(...)          → export manifest/validation
"""

# ============================================================================
# LEVEL 1: CONFIGURATION GRAMMAR
# ============================================================================

"""
cfg = jtfne.Configuration()

Fluent API:
  cfg.runtime(...)           → Configuration (with runtime settings)
  cfg.column(...)            → Configuration (with column spec)
  cfg.cell_types(...)        → Configuration (with cell type catalog)
  cfg.connectivity(...)      → Configuration (with connection rules)
  cfg.probes(...)            → Configuration (with readout probes)
  cfg.objectives(...)        → Configuration (with objective specs)
  cfg.optimizer(...)         → Configuration (with optimizer config)
  cfg.field(...)             → Configuration (with field solver)
  cfg.geometry(...)          → Configuration (with spatial geometry)

Example:
  cfg = (
      jtfne.Configuration()
      .runtime(jit=True, device="gpu")
      .column(n_neurons=1000, layer_fractions={...})
      .cell_types({"E": 0.8, "PV": 0.15, "SST": 0.05})
      .connectivity(recurrent="dense", interlaminar="sparse")
      .probes({"LFP": 128, "MUA": "all"})
  )

Every method returns a Configuration object with all prior settings preserved,
enabling inspection and modification at any stage.
"""

# ============================================================================
# LEVEL 2: MODEL GRAMMAR
# ============================================================================

"""
model = jtfne.construct(cfg)

Model methods:
  model.summary()            → dict/str (neuron counts, connectivity stats)
  model.validate()           → ValidationReport (structural checks)
  model.visualize()          → Figure (network visualization)
  model.export(path)         → save model to disk

Example:
  model = jtfne.construct(cfg)
  print(model.summary())     # {areas, layers, neurons, connectivity_matrix}
  assert model.validate()    # raises if structural issues
  model.visualize(interactive=True)
"""

# ============================================================================
# LEVEL 3: SIMULATION GRAMMAR
# ============================================================================

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

Signals methods:
  signals.V_m                → voltage array (n_trials, n_steps, n_neurons)
  signals.spikes             → spike raster (n_trials, n_steps, n_neurons)
  signals.source             → current source density (native units)
  signals.lfp_contacts       → LFP at probes (n_trials, n_steps, n_probes)
  signals.csd_contacts       → CSD at probes (n_trials, n_steps, n_probes)
  signals.field              → FieldOutput object (if solver enabled)

Signals operators:
  signals.rate(...)          → firing rates (Hz, n_trials, n_neurons or per-layer)
  signals.psd(...)           → power spectral density (freq, power)
  signals.bandpower(...)     → band power (alpha, beta, gamma, ...)
  signals.coherence(...)     → cross-frequency coherence
  signals.phase()            → instantaneous phase (per frequency)
  signals.summary()          → dict (rate, stability, firing stats)

Example:
  signals = model.simulate(seed=0)
  rate = signals.rate()
  psd = signals.psd(fmin=1, fmax=100)
  assert signals.validate()  # check for NaN/Inf
"""

# ============================================================================
# LEVEL 4: PROBE GRAMMAR
# ============================================================================

"""
TFNE doctrine:
  Emitter → Source → Field → Probe

probe = signals.probe(kind)

Probe kinds:
  "LFP-proxy"                → local field potential (Gaussian + CSD proxy)
  "EEG-proxy"                → scalp EEG (dipole + lead-field proxy)
  "MEG-proxy"                → magnetoencephalography (oriented dipole + lead-field)
  "MUA-proxy"                → multi-unit activity (spike clustering proxy)
  "CSD-proxy"                → current source density (spatial second derivative)
  "EMM"                      → extracellular monopole moment (native units)

Probe methods:
  probe.summary()            → dict (probe_type, n_contacts, contact_positions, ...)
  probe.plot()               → Figure (probe signal, spectrogram, coherence, ...)
  probe.export(path)         → save probe data (HDF5, NPZ, ...)

Example:
  lfp_probe = signals.probe("LFP-proxy")
  eeg_probe = signals.probe("EEG-proxy")
  print(lfp_probe.summary())
  lfp_probe.plot(freq_range=[1, 100])
"""

# ============================================================================
# LEVEL 5: OBJECTIVE GRAMMAR
# ============================================================================

"""
Direct objective constructors:

  jtfne.rate_targets(          → Objective
      groups={"L5": 10, "L2": 8, ...},
      weights={...}
  )

  jtfne.band_power(            → Objective
      band="gamma", target=0.2,
      layers=["L2/3"],
      weight=1.0
  )

  jtfne.phase_locking(         → Objective
      target_kappa=0.5,
      frequency_hz=40,
      weight=0.5
  )

  jtfne.csd_pattern(           → Objective
      target_pattern=array,
      weight=1.0
  )

Objective algebra (composition):

  obj_total = (
      jtfne.rate_targets({"L5": 10})
      + jtfne.band_power(band="gamma", target=0.3)
      + jtfne.phase_locking(target_kappa=0.5) * 0.5
  )

Operators:
  obj + obj                    → new Objective (weighted sum)
  obj * scalar                 → new Objective (rescale weight)
  obj / scalar                 → new Objective (rescale weight)
  -obj                         → new Objective (negate)

Objective evaluation:

  score = obj(signals)         → float (cached evaluation)
  report = obj.evaluate(signals) → ObjectiveReport (detailed breakdown)

Example:
  obj = (
      jtfne.rate_targets(
          {"L2": 10, "L5": 8, "L6": 5},
          weights={"L2": 1.0, "L5": 0.8, "L6": 0.5}
      )
      + jtfne.band_power(band="gamma", target=0.25) * 0.3
  )
  score = obj(signals)
"""

# ============================================================================
# LEVEL 6: OPTIMIZER GRAMMAR
# ============================================================================

"""
optimizer = jtfne.AGSDR(
    learning_rate=0.01,
    momentum=0.9,
    ...
)

or:

  optimizer = jtfne.GSDR(...)
  optimizer = jtfne.RandomSearch(...)

Optimizer.optimize():

  result = optimizer.optimize(
      model,
      objective,
      n_steps=1000,
      param_spec={"W_recurrent": {...}, "drive": {...}},
      early_stop_patience=50
  )

Returns OptimizationResult object.

Example:
  optimizer = jtfne.AGSDR(learning_rate=0.01)
  result = optimizer.optimize(
      model, obj,
      n_steps=500,
      param_spec={
          "drive_per_neuron": {"bounds": [-100, 100], "mask": "E_cells"},
          "W_recurrent_scale": {"bounds": [0.5, 1.5]},
      }
  )
"""

# ============================================================================
# LEVEL 7: OPTIMIZATION RESULT GRAMMAR
# ============================================================================

"""
result = optimizer.optimize(model, objective, ...)

Result attributes:
  result.best_score          → float (final objective value)
  result.best_parameters     → dict (optimized parameters)
  result.loss_trace          → array (loss over optimization steps)
  result.convergence_info    → dict (convergence stats, early_stop reason, ...)

Result methods:
  result.summary()           → str/dict (human-readable summary)
  result.plot()              → Figure (loss trace, parameter trajectory, ...)
  result.export(path)        → save result to disk
  result.apply(model)        → apply best parameters back to model
  result.validate()          → check result validity (finite, bounds, ...)

Example:
  result = optimizer.optimize(model, obj, ...)
  print(result.summary())
  result.plot(show=True)
  model_optimized = result.apply(model)
"""

# ============================================================================
# LEVEL 8: MANIFEST & VALIDATION GRAMMAR
# ============================================================================

"""
Final TFNE stage: immutable proof of computation.

  manifest = jtfne.manifest(
      cfg=cfg,
      model=model,
      signals=signals,
      objective=obj,
      result=result,
      metadata={"run_id": ..., "timestamp": ..., "tags": [...]}
  )

or:

  manifest = result.manifest(model, signals, objective)

Manifest attributes:
  manifest.config_hash       → SHA256 (canonical config fingerprint)
  manifest.model_hash        → SHA256 (compiled model state)
  manifest.signal_hashes     → {key: SHA256} (per signal array)
  manifest.metadata          → dict (run provenance, user, machine, ...)
  manifest.objective_spec    → str (human-readable objective)
  manifest.result_summary    → dict (scores, parameters, convergence)

Manifest methods:
  manifest.save(path)        → write to disk (JSON + hashes, write-once)
  manifest.validate()        → verify all hashes & integrity
  manifest.export()          → JSON-safe dict for archival

Validation:

  validation = jtfne.validation_report(
      config_valid=True,
      model_valid=True,
      signal_valid=True,
      objective_valid=True,
      result_valid=True,
      issues=[]
  )

Example:
  manifest = jtfne.manifest(
      cfg=cfg, model=model, signals=signals,
      objective=obj, result=result,
      metadata={"experiment": "Phase0_cortical_sweep", "version": "1.0"}
  )
  manifest.save("/tmp/phase0_manifest.json")
  assert manifest.validate()
"""

# ============================================================================
# FULL CHAIN EXAMPLE
# ============================================================================

"""
from pathlib import Path
import jaxfne as jtfne

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 1: Configuration (fluent, composable)
# ─────────────────────────────────────────────────────────────────────────

cfg = (
    jtfne.Configuration()
    .runtime(jit=True, x64=False, device="gpu")
    .column(
        n_neurons=1000,
        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),
        }
    )
    .cell_types({
        "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},
    })
    .connectivity(
        recurrent="dense",      # dense W_rec for local circuits
        interlaminar="sparse",  # sparse cross-layer
    )
    .probes({"LFP": 64, "MUA": "all"})
)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 2: Model (constructed, validated, inspectable)
# ─────────────────────────────────────────────────────────────────────────

model = jtfne.construct(cfg)
print(model.summary())
assert model.validate()

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 3: Signals (simulated, analyzed)
# ─────────────────────────────────────────────────────────────────────────

signals = model.simulate(seed=0, n_trials=4)
print(f"Spike rate: {signals.rate().mean():.2f} Hz")
print(f"Gamma power: {signals.bandpower('gamma').mean():.3f}")

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 4: Probe (extracted, proxy-safe)
# ─────────────────────────────────────────────────────────────────────────

lfp_probe = signals.probe("LFP-proxy")
eeg_probe = signals.probe("EEG-proxy")
print(lfp_probe.summary())

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 5: Objective (composed, weighted)
# ─────────────────────────────────────────────────────────────────────────

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

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

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 6: Optimizer (configured, executed)
# ─────────────────────────────────────────────────────────────────────────

optimizer = jtfne.AGSDR(learning_rate=0.01, momentum=0.9)
result = optimizer.optimize(
    model, objective,
    n_steps=500,
    param_spec={
        "drive_per_neuron": {
            "bounds": [-100, 100],
            "mask": "E_cells",
        },
        "W_recurrent_scale": {
            "bounds": [0.5, 1.5],
            "frozen_layers": ["L1"],
        }
    },
    early_stop_patience=50
)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 7: Result (inspected, applied)
# ─────────────────────────────────────────────────────────────────────────

print(result.summary())
result.plot(show=True)
model_optimized = result.apply(model)

# ─────────────────────────────────────────────────────────────────────────
# LEVEL 8: Manifest (proved, archived)
# ─────────────────────────────────────────────────────────────────────────

manifest = jtfne.manifest(
    cfg=cfg,
    model=model,
    signals=signals,
    objective=objective,
    result=result,
    metadata={
        "experiment": "Phase0_cortical_sweep",
        "version": "1.0",
        "user": "researcher",
        "tags": ["cortical_column", "AGSDR", "32bit"],
    }
)

manifest.save(Path("./manifests/phase0_run.json"))
assert manifest.validate()

print("✅ Full TFNE chain complete and archived")
"""

# ============================================================================
# RULE ENFORCEMENT
# ============================================================================

"""
[INVARIANT] Every jaxfne script must follow this structure:

  1. SETUP        (imports, paths, random seeds, device config)
  2. CONFIG       (fluent Configuration grammar)
  3. MODEL        (jtfne.construct)
  4. SIMULATION   (model.simulate)
  5. ANALYSIS     (signals.rate(), signals.psd(), signals.probe(...))
  6. OBJECTIVE    (jtfne.rate_targets + composition)
  7. OPTIMIZATION (optimizer.optimize)
  8. RESULT       (result.summary, result.apply)
  9. MANIFEST     (jtfne.manifest, save, validate)

Violations:
  ❌ Using raw JAX arrays instead of Signals methods
  ❌ Hardcoding layer indices instead of layer_fractions
  ❌ Direct parameter mutation instead of objectives + optimizer
  ❌ Skipping manifest/validation
  ❌ Mixing simulation and optimization in one script (separate them)

Remedy:
  - Rewrite using fluent configuration
  - Use Signals operators (.rate(), .probe(), ...)
  - Compose objectives algebraically
  - Run optimizer.optimize(), inspect result, apply to model
  - Save manifest for proof
"""
