"""Energy ledger: EP-analog vs digital training, three scenarios with uncertainty bands.
Analog core measured by SPICE (v1 column transient, integrated V*I over settle+read);
peripheral and digital terms from datasheet/literature constants (stated inline).
Scenarios: (1) MVP discrete parts, (2) T64 word-streaming, (3) integrated weight-stationary
projection. Output: pJ/MAC table + bar figure with uncertainty bands.
"""
import os
os.environ.setdefault('NGSPICE_LIBRARY_PATH', '/home/yurenh2/miniconda3/lib/libngspice.so')
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from PySpice.Spice.Netlist import Circuit, SubCircuit

N = 64
R_EQ = 1.3 * 11e3
C_BUS = 64 * 85e-12
RF, CF = 1e3, 294e-12
V_STEP = 0.1
T_SETTLE = 3.2e-6
T_READ = 1.0e-6
DWELL = T_SETTLE + T_READ

class OpAmp(SubCircuit):
    NODES = ('inp', 'inn', 'out')
    def __init__(self, name):
        super().__init__(name, *self.NODES)
        self.B('gain', 'x', self.gnd, v='1e5*(v(inp)-v(inn))')
        rp = 1e6
        cp = 1.0 / (2 * np.pi * 100.0 * rp)
        self.R('p', 'x', 'p1', rp); self.C('p', 'p1', self.gnd, cp)
        self.B('buf', 'o', self.gnd, v='v(p1)')
        self.R('out', 'o', 'out', 25)

c = Circuit('energy_col')
c.subcircuit(OpAmp('opamp'))
for i in range(N):
    c.PulseVoltageSource(f'in{i}', f'n{i}', c.gnd, initial_value=0, pulsed_value=V_STEP,
                         delay_time=0.2e-6, rise_time=50e-9, fall_time=50e-9,
                         pulse_width=1, period=2)
    c.R(f'w{i}', f'n{i}', 'sum', R_EQ)
c.C('bus', 'sum', c.gnd, C_BUS)
c.R('f', 'out', 'sum', RF); c.C('f', 'out', 'sum', CF)
c.X('amp', 'opamp', c.gnd, 'sum', 'out')
sim = c.simulator(temperature=27, nominal_temperature=27)
an = sim.transient(step_time=4e-9, end_time=0.2e-6 + DWELL)
t = np.array(an.time)
# source-delivered power: sum_i V_i * I(V_i); ngspice gives branch currents of V sources
p_src = np.zeros_like(t)
for i in range(N):
    try:
        ib = np.array(an[f'vin{i}'])          # current through source (A, into +)
    except Exception:
        ib = np.array(an[f'v.vin{i}#branch'])
    p_src += V_STEP * np.abs(ib)
E_network = float(np.trapz(p_src, t))          # J per column read (resistive + C charging)
E_per_col_net = E_network
print(f'SPICE: analog network energy per column read = {E_per_col_net*1e9:.3f} nJ '
      f'({E_per_col_net/N*1e12:.2f} pJ/MAC)')

# ---------------- ledger constants (stated assumptions) ----------------
# discrete parts:
P_OPAMP_DISC = 5e-3          # MCP6022 1 mA x 5 V quiescent, per column amp
E_ADC_DISC = 10e-9           # AD7606-class per 16-bit conversion (~100 mW / 8ch / 1 MSPS class)
E_DAC_RELOAD = 8 * 10e-12    # 8-bit latch write, ~10 pJ/bit I/O (word-streaming, per cell per use)
E_DRAM_BYTE = 20e-12         # LPDDR-class streaming, per byte
# integrated projection:
P_OPAMP_INT = 10e-6          # integrated column amp
E_ADC_INT = 0.5e-12 * 16     # ~0.5 pJ/conv-bit SAR class
# digital reference (system-level, INT8):
E_DIG_LOW, E_DIG_HIGH = 0.3e-12, 1.0e-12   # pJ/MAC incl movement, H100-class system envelope

def scenario(name, e_net_percol, p_amp, e_adc, e_stream_percell, wiring_lo=1.0, wiring_hi=3.0,
             dwell=DWELL):
    # per column-read: network + amp*dwell + one ADC conversion; per MAC = /N; plus streaming/cell
    e_core = e_net_percol * (dwell / DWELL) + p_amp * dwell + e_adc
    per_mac_lo = (e_core * wiring_lo) / N + e_stream_percell
    per_mac_hi = (e_core * wiring_hi) / N + e_stream_percell
    # EP training step = 2 settles (free+nudged) + transpose read ~ 3 column ops per MAC-use
    tr_lo, tr_hi = 3 * per_mac_lo, 3 * per_mac_hi
    # digital training step = 3x MACs (fwd+bwd) at system energy
    dig_lo, dig_hi = 3 * E_DIG_LOW, 3 * E_DIG_HIGH
    ratio_best = dig_hi / tr_lo; ratio_worst = dig_lo / tr_hi
    print(f'{name:28s} per-MAC {per_mac_lo*1e12:8.2f}-{per_mac_hi*1e12:8.2f} pJ | '
          f'train-step vs digital: {ratio_worst:6.2f}x - {ratio_best:6.2f}x '
          f'({">1 = analog wins" if ratio_best > 1 else "loses"})')
    return per_mac_lo, per_mac_hi, ratio_worst, ratio_best

print('\n=== pJ/MAC and EP-vs-digital-training energy ratio (range = wiring 1-3x + envelope) ===')
s1 = scenario('MVP discrete parts', E_per_col_net, P_OPAMP_DISC, E_ADC_DISC, 0.0)
s2 = scenario('T64 word-streaming', E_per_col_net, P_OPAMP_DISC, E_ADC_DISC,
              E_DAC_RELOAD + 1 * E_DRAM_BYTE)
# integrated: C_bus ~ 50 fF/cell -> settle ns-class; coherent dwell 100 ns (settle+read)
s3 = scenario('Integrated weight-stationary', E_per_col_net, P_OPAMP_INT, E_ADC_INT, 0.0,
              dwell=100e-9)

# ---------------- figure ----------------
names = ['MVP\n(discrete)', 'T64\n(word-stream)', 'Integrated\n(weight-stationary)']
los = [s1[0], s2[0], s3[0]]; his = [s1[1], s2[1], s3[1]]
fig, ax = plt.subplots(figsize=(8.8, 4.8))
xs = np.arange(3)
mid = [(a * b) ** 0.5 for a, b in zip(los, his)]
ax.bar(xs, [m * 1e12 for m in mid], yerr=[[(m - l) * 1e12 for m, l in zip(mid, los)],
                                          [(h - m) * 1e12 for h, m in zip(his, mid)]],
       color=['#b03a2e', '#d95f02', '#2e7d32'], alpha=0.85, capsize=6)
ax.axhspan(E_DIG_LOW * 1e12, E_DIG_HIGH * 1e12, color='#2c6fbb', alpha=0.18)
ax.text(2.35, E_DIG_HIGH * 1e12 * 1.1, 'digital INT8 system\n0.3–1 pJ/MAC', fontsize=8.5,
        color='#2c6fbb', ha='right')
ax.set_yscale('log')
ax.set_xticks(xs); ax.set_xticklabels(names)
ax.set_ylabel('pJ per MAC (log)')
ax.set_title('EP-analog energy per MAC — SPICE-measured core + datasheet periphery\n'
             '(bands: schematic-vs-layout wiring 1–3×)')
fig.tight_layout()
fig.savefig('/home/yurenh2/ept/assets/figs/fig_energy_ledger.png', dpi=150)
print('DONE_ENERGY')