"""
Tesla Coil Spark Course - Circuit Diagram Generation

Generates circuit schematics using schemdraw.
Run from spark-lessons directory.

Usage: python generate_circuits.py
"""

import schemdraw
import schemdraw.elements as elm
from pathlib import Path
import matplotlib.pyplot as plt

# Directories
BASE_DIR = Path(__file__).parent
ASSETS_DIRS = {
    'fundamentals': BASE_DIR / 'lessons' / '01-fundamentals' / 'assets',
    'optimization': BASE_DIR / 'lessons' / '02-optimization' / 'assets',
    'spark-physics': BASE_DIR / 'lessons' / '03-spark-physics' / 'assets',
    'advanced-modeling': BASE_DIR / 'lessons' / '04-advanced-modeling' / 'assets',
    'shared': BASE_DIR / 'assets' / 'shared',
}

def save_circuit(drawing, filename, directory='fundamentals'):
    """Save circuit diagram"""
    filepath = ASSETS_DIRS[directory] / filename
    drawing.save(str(filepath), dpi=150)
    print(f"[OK] Generated: {filepath}")


# ============================================================================
# PART 1: FUNDAMENTALS CIRCUITS
# ============================================================================

def generate_geometry_to_circuit():
    """Image 2: Geometry to circuit schematic translation"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=12, font='sans-serif')

        # Draw the circuit on the right side
        # Topload node at top
        d += elm.Line().right(1).label('Topload', loc='top')
        d.push()

        # Parallel R and C_mut
        d += elm.Line().down(0.5)
        d.push()
        d += elm.Resistor().down(1.5).label('R', loc='right')
        d.pop()
        d += elm.Capacitor().down(1.5).label('C_mut', loc='right').at((1, d.here[1]))
        d += elm.Line().left(1)

        # Series point (spark tip node)
        d += elm.Dot().label('Spark Tip', loc='right', ofst=0.3)
        d += elm.Line().down(0.5)

        # C_sh to ground
        d += elm.Capacitor().down(1.5).label('C_sh', loc='right')
        d += elm.Ground()

        # Title is implicit in context - no annotation needed

    save_circuit(d, 'geometry-to-circuit.png', 'fundamentals')


def generate_current_paths_diagram():
    """Image 6: Tesla coil showing all current paths"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=10, font='sans-serif')

        # Primary circuit (left)
        d += elm.SourceSin().label('Drive')
        d += elm.Capacitor().right(1.5).label('C_pri')
        d += elm.Inductor().down(2).label('L_pri')
        d += elm.Line().left(1.5)
        d += elm.Ground()

        # Coupling to secondary
        d.move(2, 1.5)
        d += elm.Inductor().up(3).label('L_sec', loc='right')
        d.push()

        # Topload capacitance
        d += elm.Line().right(0.5)
        d += elm.Capacitor().right(1).label('C_top')
        d += elm.Line().down(0.5)

        # Spark circuit
        d.push()
        d += elm.Capacitor().down(1).label('C_mut', loc='right')
        d += elm.Line().down(0.5)
        d += elm.Capacitor().down(1).label('C_sh', loc='right')
        d += elm.Ground()
        d.pop()

        # Ground path
        d += elm.Line().right(1.5)
        d += elm.Ground()

        # Add current labels
        d.here = (0, -2.5)
        d += elm.Annotate().label('I_base', fontsize=10, color='red')

    save_circuit(d, 'current-paths-diagram.png', 'fundamentals')


# ============================================================================
# PART 2: OPTIMIZATION CIRCUITS
# ============================================================================

def generate_thevenin_equivalent_circuit():
    """Image 12: Thévenin equivalent with spark load"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=12, font='sans-serif')

        # Thévenin source
        d += elm.SourceV().label('V_th')
        d.push()

        # Z_th (impedance)
        d += elm.Resistor().right(1.5).label('R_th')
        d += elm.Capacitor().right(1.5).label('X_th', loc='bottom')

        # Connection point
        d += elm.Dot()
        d.push()

        # Load (spark)
        d += elm.Line().down(0.5)
        d += elm.Resistor().down(1.5).label('R_spark', loc='right')
        d += elm.Capacitor().down(1.5).label('X_spark', loc='right')
        d += elm.Ground()

        # Close circuit
        d.pop()
        d += elm.Line().down(4.5)
        d += elm.Line().left(3)

        # Add formula annotation
        d.here = (1, -5.5)
        d += elm.Annotate(ofst=(0, -0.5)).label(
            'P = 0.5|V_th|² Re{Z_spark} / |Z_th+Z_spark|²',
            fontsize=11
        )

    save_circuit(d, 'thevenin-equivalent-circuit.png', 'optimization')


# ============================================================================
# PART 3: SPARK PHYSICS CIRCUITS
# ============================================================================

def generate_capacitive_divider_circuit():
    """Image 25: Capacitive divider circuit"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=12, font='sans-serif')

        # Voltage source (topload)
        d += elm.Line().right(1).label('V_topload', loc='top')
        d += elm.Dot()
        d.push()

        # Parallel R and C_mut
        d += elm.Line().down(0.5)
        d.push()
        d += elm.Resistor().down(1.5).label('R')
        d.pop()
        d += elm.Capacitor().right(1.5).down(1.5).label('C_mut')
        d += elm.Line().left(1.5)

        # V_tip measurement point
        d += elm.Dot().label('V_tip', loc='right', ofst=0.3)
        d += elm.Line().down(0.5)

        # C_sh to ground
        d += elm.Capacitor().down(1.5).label('C_sh = L×6.6pF/m', loc='right')
        d += elm.Ground()

        # Add formula
        d.here = (0, -5)
        d += elm.Annotate().label(
            'V_tip = V_topload × C_mut/(C_mut + C_sh)',
            fontsize=11
        )

    save_circuit(d, 'capacitive-divider-circuit.png', 'spark-physics')


# ============================================================================
# PART 4: ADVANCED MODELING CIRCUITS
# ============================================================================

def generate_lumped_model_schematic():
    """Image 28: Lumped model circuit schematic"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=11, font='sans-serif')

        # Topload connection
        d += elm.Line().right(1).label('Topload', loc='top')
        d += elm.Dot().label('Port')
        d.push()

        # Parallel combination
        d += elm.Line().down(0.3)
        d.push()

        # R branch
        d += elm.Resistor().down(2).label('R', loc='left')

        # C_mut branch
        d.pop()
        d += elm.Capacitor().right(2).down(2).label('C_mut', loc='right')
        d += elm.Line().left(2)

        # Spark tip node
        d += elm.Dot().label('Spark Tip', loc='right', ofst=0.3)

        # C_sh to ground
        d += elm.Line().down(0.3)
        d += elm.Capacitor().down(1.5).label('C_sh', loc='right')
        d += elm.Ground()

        # Add typical values
        d.here = (0, -5)
        d += elm.Annotate().label(
            'Typical: R=50kΩ, C_mut=8pF, C_sh=6pF',
            fontsize=10
        )

    save_circuit(d, 'lumped-model-schematic.png', 'advanced-modeling')


def generate_distributed_model_structure():
    """Image 32: nth-order distributed model structure"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=9, font='sans-serif')

        # Topload
        d += elm.Line().right(0.5).label('Topload', loc='top')
        d += elm.Dot().label('Node 0')

        # Segment 1
        d.push()
        d += elm.Capacitor().down(1.2).label('C_01', loc='left', ofst=-0.2)
        d += elm.Dot().label('Node 1', loc='right', ofst=0.2)
        d.push()
        d += elm.Resistor().right(1.5).label('R_1', loc='top')
        d.pop()
        d += elm.Capacitor().down(1.2).label('C_1,gnd', loc='left')
        d += elm.Ground()

        # Segment 2
        d.pop()
        d += elm.Line().right(3)
        d.push()
        d += elm.Capacitor().down(1.2).label('C_12', loc='left', ofst=-0.2)
        d += elm.Dot().label('Node 2', loc='right', ofst=0.2)
        d.push()
        d += elm.Resistor().right(1.5).label('R_2', loc='top')
        d.pop()
        d += elm.Capacitor().down(1.2).label('C_2,gnd', loc='left')
        d += elm.Ground()

        # Ellipsis
        d.pop()
        d += elm.Line().right(1.5)
        d += elm.Dot()
        d += elm.Line().right(0.3).linestyle('dotted')
        d += elm.Line().right(0.3)
        d += elm.Dot()
        d += elm.Line().right(1.5)

        # Segment n
        d.push()
        d += elm.Capacitor().down(1.2).label('C_n-1,n', loc='left', ofst=-0.2)
        d += elm.Dot().label('Node n', loc='right', ofst=0.2)
        d.push()
        d += elm.Resistor().right(1.5).label('R_n', loc='top')
        d.pop()
        d += elm.Capacitor().down(1.2).label('C_n,gnd', loc='left')
        d += elm.Ground()

        # Add note
        d.here = (3, -3.5)
        d += elm.Annotate().label(
            'n = 5-20 segments\n(n+1)×(n+1) capacitance matrix',
            fontsize=9
        )

    save_circuit(d, 'distributed-model-structure.png', 'advanced-modeling')


# ============================================================================
# SHARED CIRCUITS
# ============================================================================

def generate_tesla_coil_system_overview():
    """Image 44: Complete Tesla coil system diagram"""
    with schemdraw.Drawing(show=False) as d:
        d.config(fontsize=10, font='sans-serif')

        # Primary side
        d += elm.SourceSin().label('Drive\nSource')
        d += elm.Line().right(0.5)
        d += elm.Switch().label('IGBT/FET')
        d += elm.Line().right(0.5)
        d += elm.Capacitor().right(1.5).label('MMC\n(C_pri)')
        d += elm.Inductor().down(3).label('L_primary', loc='bottom')
        d += elm.Line().left(3.5)
        d += elm.Ground()

        # Secondary side (coupled)
        d.move(4, 2)
        d += elm.Inductor().up(4).label('L_secondary', loc='right')
        d += elm.Line().up(0.5)

        # Topload
        d += elm.Capacitor().right(1.5).label('C_topload')
        d.push()

        # Spark
        d += elm.Line().down(1)
        d += elm.Gap().down(2).label('Spark\nGap')
        d += elm.Line().down(1)
        d += elm.Ground().label('Strike\nPoint')

        # Ground return
        d.pop()
        d += elm.Line().right(2)
        d += elm.Line().down(5.5)
        d += elm.Ground()

        # Add coupling annotation
        d.here = (2, 0)
        d += elm.Annotate(ofst=(0, 2)).label('k = 0.1-0.2', fontsize=10)

        # Add title
        d.here = (0, 7)
        d += elm.Annotate().label(
            'Double-Resonant Solid State Tesla Coil (DRSSTC)',
            fontsize=12
        )

    save_circuit(d, 'tesla-coil-system-overview.png', 'shared')


# ============================================================================
# MAIN
# ============================================================================

def main():
    print("\n" + "="*60)
    print("TESLA COIL SPARK COURSE - CIRCUIT DIAGRAM GENERATION")
    print("="*60)

    print("\nGenerating Part 1 circuits...")
    generate_geometry_to_circuit()
    generate_current_paths_diagram()

    print("\nGenerating Part 2 circuits...")
    generate_thevenin_equivalent_circuit()

    print("\nGenerating Part 3 circuits...")
    generate_capacitive_divider_circuit()

    print("\nGenerating Part 4 circuits...")
    generate_lumped_model_schematic()
    generate_distributed_model_structure()

    print("\nGenerating shared circuits...")
    generate_tesla_coil_system_overview()

    print("\n" + "="*60)
    print("CIRCUIT GENERATION COMPLETE!")
    print("="*60)
    print(f"\nTotal circuit diagrams generated: 7")
    print("="*60 + "\n")


if __name__ == '__main__':
    main()