#!/usr/bin/env python3
"""
jaxfne-cortical-column-default — the canonical 1K-neuron reference template

Use this as the default cortical column architecture for all jaxfne work
unless explicitly specified otherwise (e.g., "use 10-layer thalamocortical model").

GROUND TRUTH (verified 2026-06-17 via construct().neuron_table()). The two laws:
  1. E peaks DEEP   — excitatory FRACTION rises monotonically with depth (L6 ≈ 90% E).
  2. I peaks SUPERFICIAL — inhibitory FRACTION highest superficial (L1 50% I); the
     largest inhibitory NEURON COUNT sits in the dense superficial L2 / L2-3.
Overall realistic E:I ≈ 77:23 (NOT the old over-inhibitory 63:37 / 41:59).

2026-06-19 INTEGRATION — build_laminar_column(ei_profile="canonical") now applies
canonical biophysics natively at construct time (core._apply_canonical_biophysics):
  - Per-layer fractions = user reference FRACS_LAYER (more inhibitory superficial:
    L1/L2/L3 are 50% I; L1 VIP-rich 0.35; L4 PV 0.20; L6 0.95 E, no PV/VIP; ~72E:28I).
  - SST a=0.05 (intermediate E=0.02 < SST < PV=0.10).
  - Deep-E "larger": source_scale 1.0->1.8, a 0.020->0.015, d 8->10 graded by depth Z.
  - PV<->E local connectivity x3, distance-gated (PING feedback-inhibition loop).
  - Random v0 ~ U(-70,0) ALWAYS (every construct; opt-out cfg.runtime(random_v0=False))
    -- removes onset-response synchrony bias.
  - Homeostasis size-aware: homeostasis_params={"k_gain_size_scaled":True} -> per-neuron
    k_gain proportional to 1/source_scale (larger/deep-E taxed less).
Effect: synchrony kappa drops ~0.040->0.016 at the same ~10 Hz mean rate.

This template embodies:
  - Biologically plausible L2/3 dominance (45% of neurons)
  - Realistic layer thickness & neuron density (L2/3 densest)
  - Realistic E:I gradient: superficial inhibitory-rich, deep excitatory-dominant
  - Normalized depth coordinates (scale-agnostic)
  - Tractable compute (1K neurons; construct ~2s, simulate 1000ms ~1s on CPU JAX)
  - Scalable to 10K for full-resolution studies
"""

# ============================================================================
# LAYER ARCHITECTURE (1000 neurons)
# ============================================================================

"""
| Layer | Thickness | Neurons | I-fraction | role                    |
|-------|----------:|--------:|-----------:|-------------------------|
| L1    |      0.10 |     100 |       50%  | interneuron-rich (no PV)|
| L2    |      0.15 |     250 |       30%  | I-COUNT peak (dense)    |  ← densest
| L3    |      0.15 |     200 |       25%  | mixed                   |  ← densest
| L4    |      0.10 |     100 |       20%  | PV feedforward input    |
| L5    |      0.30 |     200 |       12%  | E projection neurons    |
| L6    |      0.20 |     150 |       10%  | E corticothalamic       |  ← E-frac peak
| Total |      1.00 |    1000 |    ~23%    | 77E : 23I               |

Key properties:
  - E-fraction RISES with depth: L1 50% -> L2 70% -> L3 75% -> L4 80% -> L5 88% -> L6 90%.
  - I-fraction FALLS with depth: 50 -> 30 -> 25 -> 20 -> 12 -> 10%.
  - I-COUNT peak is L2 (~75 interneurons) — the dense superficial layer holds the
    biggest inhibitory pool even though L1 has the highest I-fraction.
  - L2/3 (superficial, output): 45% of all neurons (450/1000).
  - L5 thickness 0.30 (thickest); L5+L6 (deep) = 35% of neurons.
  - PV concentrates in L4 (fast feedforward inhibition); L1 has NO PV (VIP/SST only).
"""

# ============================================================================
# JAXFNE LAYER FRACTIONS (normalized depth intervals; WIDTH ∝ neuron count)
# ============================================================================

"""
LAYER_FRACTIONS = {
    "L1": (0.0, 0.1),    # top 10%
    "L2": (0.1, 0.35),   # next 25%
    "L3": (0.35, 0.55),  # next 20%
    "L4": (0.55, 0.65),  # next 10%
    "L5": (0.65, 0.85),  # next 20%
    "L6": (0.85, 1.0),   # bottom 15%
}
Band WIDTH is proportional to neuron count (count ∝ width): 100,250,200,100,200,150.
  - L1 superficial (cortical surface at depth 0); L6 deep (white matter at depth 1).
"""

# ============================================================================
# CELL TYPE DISTRIBUTION (by layer) — GROUND TRUTH gradient
# ============================================================================

"""
LAYER_CELL_TYPE_FRAC = {
    "L1": {"E": 0.50, "PV": 0.00, "SST": 0.15, "VIP": 0.35},  # 50% I (no PV; VIP/SST)
    "L2": {"E": 0.70, "PV": 0.15, "SST": 0.10, "VIP": 0.05},  # 30% I (I-count peak)
    "L3": {"E": 0.75, "PV": 0.13, "SST": 0.08, "VIP": 0.04},  # 25% I
    "L4": {"E": 0.80, "PV": 0.12, "SST": 0.05, "VIP": 0.03},  # 20% I (PV feedforward)
    "L5": {"E": 0.88, "PV": 0.06, "SST": 0.04, "VIP": 0.02},  # 12% I (E projection)
    "L6": {"E": 0.90, "PV": 0.05, "SST": 0.03, "VIP": 0.02},  # 10% I (E-frac peak)
}

Principle:
  - Superficial (L1–L3): inhibitory-RICH (50% -> 25% I), modulation/integration circuits.
  - Deep (L5–L6): strongly excitatory (88–90% E), projection-neuron dominant.
  - PV (fast inhibition) peaks at L4 thalamorecipient layer; absent in L1.
  - Matches biology: superficial layers integrate & modulate, deep layers project.
"""

# ============================================================================
# QUICKSTART: MINIMAL JAXFNE SETUP (fluent grammar, v0.4.0 — verified)
# ============================================================================

"""
import jaxfne as jtfne
import collections
jtfne.enable_x64()

LAYERS = ["L1","L2","L3","L4","L5","L6"]

# z-bands: WIDTH ∝ neuron count
ZBANDS = {
    "L1": (0.00, 0.10), "L2": (0.10, 0.35), "L3": (0.35, 0.55),
    "L4": (0.55, 0.65), "L5": (0.65, 0.85), "L6": (0.85, 1.00),
}

# Cell-type composition per layer (GROUND TRUTH gradient: E-deep, I-superficial)
LAYER_CELL_TYPE_FRAC = {
    "L1": {"E": 0.50, "PV": 0.00, "SST": 0.15, "VIP": 0.35},
    "L2": {"E": 0.70, "PV": 0.15, "SST": 0.10, "VIP": 0.05},
    "L3": {"E": 0.75, "PV": 0.13, "SST": 0.08, "VIP": 0.04},
    "L4": {"E": 0.80, "PV": 0.12, "SST": 0.05, "VIP": 0.03},
    "L5": {"E": 0.88, "PV": 0.06, "SST": 0.04, "VIP": 0.02},
    "L6": {"E": 0.90, "PV": 0.05, "SST": 0.03, "VIP": 0.02},
}

cfg = (jtfne.laminar_cortex_config(
            seed=0, duration_ms=1000.0, dt_ms=0.5, areas=["V1"], layers=LAYERS,
            n=1000, emitter="izhikevich",
            baseline_drive_by_cell_type={"E":5.0,"PV":5.0,"SST":5.0,"VIP":5.0})
       .layer_fractions(layer_fractions=ZBANDS)          # WIDTH ∝ count -> per-layer counts
       .area_layer_cell_types("V1", LAYER_CELL_TYPE_FRAC))  # per-layer E/I composition
model = jtfne.construct(cfg)

# Verify (neuron_table() returns a LIST of dict rows: neuron_id, area, layer, cell_type, x,y,z)
lc = collections.Counter((r["layer"], r["cell_type"]) for r in model.neuron_table())
# L1: 50E/0PV/15SST/35VIP (100)   L2: 175E/.. (250)   L3: 150E/.. (200)
# L4: 80E/.. (100)   L5: 176E/.. (200)   L6: 135E/.. (150)   TOTAL 77E:23I

WARNING: laminar_cortex_config(cell_types={"E":0.8,...}) GLOBAL weights produce the
WRONG per-layer gradient (deep over-inhibitory, 41:59). ALWAYS override per layer with
.area_layer_cell_types(area, {...}).

NOTE: layer_celltype_count_table / column_density_table raise NotImplementedError on a
Configuration — read counts from construct(cfg).neuron_table() instead.
"""

# ============================================================================
# DRIVE / SANITY CALIBRATION (verified Step 3, 2026-06-17)
# ============================================================================

"""
baseline_drive_by_cell_type = 5.0  -> mean rate ≈ 18 Hz (in 8–25 Hz band) at n=1000.
  Drive sweep: 4 -> 12.8 Hz, 6 -> 23.4 Hz, 8 -> 32.6 Hz (hot), 0 -> silent.
Sanity gate: Vm rest ≈ −67 mV, spike peak +30 mV (hard reset), all finite.
  construct ~1.8 s, simulate(1000ms) ~1 s on CPU.
Proxy fields auto-computed (linear_solver, NO PDE): lfp_proxy, csd_proxy, source_proxy,
  each (n_steps, n_contacts), finite. eeg_proxy/meg_proxy NOT auto (need explicit lead-field).
Signal keys are *_proxy (NO *_like, NO *_contacts).
"""

# ============================================================================
# STIMULUS DESIGN: LAYER-SPECIFIC DC DRIVE
# ============================================================================

"""
Layer-specific current injection via the neuron table (list of dict rows):

import numpy as np
nt = model.neuron_table()
drive_per_neuron = np.zeros(len(nt), dtype=np.float64)
for i, r in enumerate(nt):
    if r["layer"] in ("L2", "L3") and r["cell_type"] == "E":
        drive_per_neuron[i] = 10.0
model = model.with_emitter_parameters(drive_per_neuron=drive_per_neuron)
"""

# ============================================================================
# SPECTROLAMINAR ANALYSIS
# ============================================================================

"""
signals = jtfne.simulate(model, jtfne.simulation(duration_ms=1000.0, dt_ms=0.5, seed=0))
fig = jtfne.vis.spectrolaminar_suite(signals)   # preferred laminar readout

# Spectrolaminar = depth × FREQUENCY relative power:
#   alpha-beta deep, gamma superficial (the band-crossover is SCALE-EMERGENT, present at
#   ~10k/1000ms, absent at n<=4000). Do not conflate with field-laminar (depth × TIME).
"""

# ============================================================================
# OVERRIDE PROTOCOL
# ============================================================================

"""
To use a DIFFERENT architecture:
  1. State explicitly ("uniform 200 neurons/layer", "balanced 50/50 E/I", "10-layer model").
  2. Provide full layer spec (thicknesses, counts, cell types).
  3. If not stated, this canonical 1K template (77E:23I, E-deep / I-superficial) is the prior.
"""
