#!/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").

This template embodies:
  - Biologically plausible L2/3 dominance (45% of neurons)
  - Realistic layer thickness & neuron density (L2/3 densest)
  - Normalized depth coordinates (scale-agnostic)
  - Balanced superficial/deep coverage
  - Tractable compute (1K neurons ≈ 120–150s/trial on CPU)
  - Scalable to 10K for full-resolution studies
"""

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

"""
| Layer | Thickness | Neurons | Density (neurons/unit) |
|-------|----------:|--------:|---------------------:|
| L1    |      0.10 |     100 |          1000 /unit  |
| L2    |      0.15 |     250 |          1667 /unit  |  ← densest
| L3    |      0.15 |     200 |          1333 /unit  |  ← densest
| L4    |      0.10 |     100 |          1000 /unit  |
| L5    |      0.30 |     200 |           667 /unit  |
| L6    |      0.20 |     150 |           750 /unit  |
|-------|-----------|---------|
| Total |      1.00 |    1000 |                      |

Key properties:
  - L2/3 (superficial, output): 45% of all neurons (450/1000)
  - L2 peak density: 1667 neurons/unit (highest)
  - L5 thickness: 0.30 (thickest layer, prominent in motor/sensory)
  - L5+L6 (deep): 35% of neurons (350/1000)
  - Normalized depth [0, 1] enables layer-specific stimulus masking

Use case: spectrolaminar analysis, canonical microcircuit, hierarchical models
"""

# ============================================================================
# JAXFNE LAYER FRACTIONS (normalized depth intervals)
# ============================================================================

"""
Python dict for jaxfne.Configuration:

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 30%
    "L6": (0.85, 1.0),   # bottom 15%
}

Depth interpretation:
  - L1 is superficial (top, cortical surface at depth 0)
  - L6 is deep (bottom, approaching white matter at depth 1)
  - All depths are relative to column thickness (arbitrary units: μm, mm, etc.)
"""

# ============================================================================
# CELL TYPE DISTRIBUTION (by layer)
# ============================================================================

"""
LAYER_CELL_TYPE_FRAC = {
    "L1": {"E": 0.35, "PV": 0.30, "SST": 0.20, "VIP": 0.15},  # inhibitory-heavy
    "L2": {"E": 0.45, "PV": 0.25, "SST": 0.18, "VIP": 0.12},  # mixed
    "L3": {"E": 0.55, "PV": 0.20, "SST": 0.15, "VIP": 0.10},  # E-biased
    "L4": {"E": 0.65, "PV": 0.18, "SST": 0.12, "VIP": 0.05},  # E-dominant (thalamic input)
    "L5": {"E": 0.85, "PV": 0.08, "SST": 0.05, "VIP": 0.02},  # strongly E (projection neurons)
    "L6": {"E": 0.90, "PV": 0.06, "SST": 0.03, "VIP": 0.01},  # strongly E (corticothalamic)
}

Principle:
  - Superficial (L1–L4): More inhibitory (PV/SST/VIP), mixed E/I circuits
  - Deep (L5–L6): Strongly excitatory, projection-neuron dominant
  - Matches biology: superficial layers integrate & modulate, deep layers project
"""

# ============================================================================
# QUICKSTART: MINIMAL JAXFNE SETUP
# ============================================================================

"""
import jaxfne as jtfne
from jaxfne import tutorial_utils as tu

# Layer architecture (normalized depths)
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 type composition (E/I ratio per layer)
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},
}

# Build configuration
layer_fracs_tuple = 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",),                          # single area
    layers=("L1", "L2", "L3", "L4", "L5", "L6"),
    layer_fractions=layer_fracs_tuple,      # canonical normalized depths
    layer_cell_type_frac=LAYER_CELL_TYPE_FRAC,
    n_neuron_per_column=1000,               # 1K neurons
    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)
neurons_df = model["neurons"]

# Verify layer counts
for layer in ("L1", "L2", "L3", "L4", "L5", "L6"):
    count = len(neurons_df[neurons_df["layer"] == layer])
    print(f"{layer}: {count} neurons")

# Output:
# L1: 100 neurons
# L2: 250 neurons
# L3: 200 neurons
# L4: 100 neurons
# L5: 200 neurons
# L6: 150 neurons
"""

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

"""
Using normalized depths [0, 1] to apply layer-specific current injection:

# Apply 10 nA to L2/L3 E cells only
drive_per_neuron = np.zeros(len(neurons_df), dtype=np.float32)

# L2 + L3 spans depth [0.1, 0.55]
layer_mask = (
    (neurons_df["layer"].isin(["L2", "L3"]))
    & (neurons_df["cell_type"] == "E")
)
drive_per_neuron[layer_mask.values] = 10.0

# Time-gating: apply drive only during [t_on, t_off] ms
n_steps = int(duration_ms / dt_ms)
t_indices = np.arange(n_steps, dtype=np.float32) * dt_ms
time_gate = (t_indices >= t_on) & (t_indices < t_off)
stimulus = np.ones((n_steps, len(neurons_df)), dtype=np.float32) * drive_per_neuron[np.newaxis, :]
stimulus[~time_gate] = 0

# Simulate with layer-specific, time-gated drive
signals = tu.simulate_laminar_trials(model, cfg, stimulus=stimulus, n_trials=4)
"""

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

"""
The default template is optimized for spectrolaminar decomposition:

signals = model.simulate(seed=0)
scores, specs = tu.summarize_spectrolaminar_similarity(signals, cfg)

# Spectrolaminar output: depth × frequency relative power
# - L2/3 dominance at alpha-beta (~10–30 Hz)
# - L5/L6 gamma peaks (~40–100 Hz)
# - Depth-dependent phase lags (thalamic→L4→L2/3→L5/6 routing)

from jaxfne.vis.tutorial_panels import spectrolaminar_suite_3panel

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

# Output: 3-panel figure (LFP, CSD, spectrolaminar heatmap)
"""

# ============================================================================
# HIERARCHICAL MULTI-AREA STACKING
# ============================================================================

"""
The normalized depth framework enables layer-matched stacking of multiple areas:

AREAS = ["V1", "V2", "V4", "PFC"]

# Each area uses the same layer_fractions, layer_cell_type_frac
# Different area properties (connectivity, recurrent strength) can be applied per-area

cfg = (
    jtfne.Configuration()
    .runtime(...)
    .multi_area(areas=AREAS)
    .area_connectivity(
        V1_to_V2="feedforward_L5_to_L1",
        V2_to_V4="feedforward_L5_to_L1",
        V4_to_PFC="feedforward_L5_to_L1",
    )
    .area_layer_fractions(LAYER_FRACTIONS)  # shared across all areas
    .area_cell_types(LAYER_CELL_TYPE_FRAC)  # shared across all areas
)

model = jtfne.construct(cfg)
signals = model.simulate(seed=0)

# Spectrolaminar analysis per area
for area in AREAS:
    area_signals = signals.extract(area)
    scores, specs = tu.summarize_spectrolaminar_similarity(area_signals, cfg)
    # ... plot per-area spectrolaminar profiles
"""

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

"""
To use a DIFFERENT architecture in discussions:
  1. State explicitly: "using 10-layer thalamocortical model" or "uniform 200 neurons/layer"
  2. Provide full layer spec (thicknesses, neuron counts, cell types)
  3. If not stated, assume the canonical 1K template above

This ensures reproducibility and prevents silent architectural drift.
"""

# ============================================================================
# PROPERTIES OF THE DEFAULT TEMPLATE
# ============================================================================

"""
Why this template?

✅ Biologically plausible
  - L2/3 superficial dominance (45%) matches cortical output bottleneck
  - L5 thickness (0.30) reflects deep-layer projection neuron expansion
  - E/I ratios match observed coronal anatomy

✅ Computationally efficient
  - 1K neurons: ~120–150s/trial on CPU JAX
  - Scales linearly to 10K for full-resolution studies
  - Fast enough for parameter sweeps (Phase 0–2)

✅ Analytically clear
  - Spectrolaminar analysis shows L2/3 vs L5/6 frequency segregation
  - Layer-specific stimulation via normalized depths
  - Canonical reference for comparing across studies

✅ Extensible
  - Layer fractions already normalized → multi-area stacking
  - Cell type fracs are per-layer parameters → easy to adjust E/I balance
  - No hardcoded absolute depths → works at any physical scale

Use this as the default prior. If you need something else, say so explicitly.
"""
