#!/usr/bin/env python3
"""
EdgeFEM Geometry Generators

Python-based geometry generation for common antenna and array structures.
Outputs Gmsh .geo files that can be meshed with `gmsh -3 file.geo -o file.msh`
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple
import math


@dataclass
class RectWaveguide:
    """Rectangular waveguide geometry parameters."""
    a: float = 0.02286  # Width (m) - WR-90 default
    b: float = 0.01016  # Height (m) - WR-90 default
    length: float = 0.05  # Length (m)
    mesh_size: float = None  # Auto-calculated if None

    def __post_init__(self):
        if self.mesh_size is None:
            self.mesh_size = min(self.a, self.b) / 5

    def generate(self, path: str):
        """Generate Gmsh .geo file for rectangular waveguide."""
        geo = f"""// Rectangular Waveguide
// Generated by EdgeFEM geometry generator
// a = {self.a*1000:.3f} mm, b = {self.b*1000:.3f} mm, L = {self.length*1000:.3f} mm

lc = {self.mesh_size};

// Cross-section at z=0
Point(1) = {{0, 0, 0, lc}};
Point(2) = {{{self.a}, 0, 0, lc}};
Point(3) = {{{self.a}, {self.b}, 0, lc}};
Point(4) = {{0, {self.b}, 0, lc}};

Line(1) = {{1, 2}};
Line(2) = {{2, 3}};
Line(3) = {{3, 4}};
Line(4) = {{4, 1}};

Curve Loop(1) = {{1, 2, 3, 4}};
Plane Surface(1) = {{1}};

// Extrude to create waveguide
Extrude {{0, 0, {self.length}}} {{ Surface{{1}}; }}

// Physical groups
Physical Surface("PEC", 1) = {{13, 17, 21, 25}};  // Walls
Physical Surface("Port1", 2) = {{1}};              // Input (z=0)
Physical Surface("Port2", 3) = {{26}};             // Output (z=L)
Physical Volume("Air", 100) = {{1}};

// Mesh settings
Mesh.Algorithm3D = 4;  // Frontal Delaunay
Mesh.Optimize = 1;
"""
        Path(path).write_text(geo)
        print(f"Generated: {path}")
        return path


@dataclass
class PatchAntenna:
    """Microstrip patch antenna geometry parameters."""
    patch_width: float = 0.038  # Patch width (m)
    patch_length: float = 0.029  # Patch length (m)
    substrate_height: float = 0.00157  # Substrate thickness (m)
    substrate_eps_r: float = 4.4  # Relative permittivity (FR4)
    ground_extend: float = 0.02  # Ground plane extension beyond patch
    air_height: float = 0.05  # Air box height above patch
    feed_inset: float = 0.005  # Feed point inset from edge
    feed_width: float = 0.003  # Microstrip feed width
    mesh_size: float = None

    def __post_init__(self):
        if self.mesh_size is None:
            self.mesh_size = min(self.patch_width, self.patch_length) / 8

    def generate(self, path: str):
        """Generate Gmsh .geo file for patch antenna."""
        gw = self.patch_width + 2 * self.ground_extend
        gl = self.patch_length + 2 * self.ground_extend
        px = self.ground_extend
        py = self.ground_extend

        geo = f"""// Microstrip Patch Antenna
// Generated by EdgeFEM geometry generator
// Patch: {self.patch_width*1000:.2f} x {self.patch_length*1000:.2f} mm
// Substrate: h={self.substrate_height*1000:.3f} mm, eps_r={self.substrate_eps_r}

lc = {self.mesh_size};
lc_fine = {self.mesh_size/2};

// Ground plane corners (z=0)
Point(1) = {{0, 0, 0, lc}};
Point(2) = {{{gw}, 0, 0, lc}};
Point(3) = {{{gw}, {gl}, 0, lc}};
Point(4) = {{0, {gl}, 0, lc}};

Line(1) = {{1, 2}};
Line(2) = {{2, 3}};
Line(3) = {{3, 4}};
Line(4) = {{4, 1}};

Curve Loop(1) = {{1, 2, 3, 4}};
Plane Surface(1) = {{1}};

// Substrate top (z=h)
Point(5) = {{0, 0, {self.substrate_height}, lc}};
Point(6) = {{{gw}, 0, {self.substrate_height}, lc}};
Point(7) = {{{gw}, {gl}, {self.substrate_height}, lc}};
Point(8) = {{0, {gl}, {self.substrate_height}, lc}};

// Patch corners on substrate
Point(9) = {{{px}, {py}, {self.substrate_height}, lc_fine}};
Point(10) = {{{px + self.patch_width}, {py}, {self.substrate_height}, lc_fine}};
Point(11) = {{{px + self.patch_width}, {py + self.patch_length}, {self.substrate_height}, lc_fine}};
Point(12) = {{{px}, {py + self.patch_length}, {self.substrate_height}, lc_fine}};

// Air box top
Point(13) = {{0, 0, {self.substrate_height + self.air_height}, lc}};
Point(14) = {{{gw}, 0, {self.substrate_height + self.air_height}, lc}};
Point(15) = {{{gw}, {gl}, {self.substrate_height + self.air_height}, lc}};
Point(16) = {{0, {gl}, {self.substrate_height + self.air_height}, lc}};

// Build substrate top with patch hole
Line(5) = {{5, 6}};
Line(6) = {{6, 7}};
Line(7) = {{7, 8}};
Line(8) = {{8, 5}};

Line(9) = {{9, 10}};
Line(10) = {{10, 11}};
Line(11) = {{11, 12}};
Line(12) = {{12, 9}};

Curve Loop(2) = {{5, 6, 7, 8}};
Curve Loop(3) = {{9, 10, 11, 12}};
Plane Surface(2) = {{2, 3}};  // Substrate top (with hole)
Plane Surface(3) = {{3}};      // Patch

// Substrate sides
Line(13) = {{1, 5}};
Line(14) = {{2, 6}};
Line(15) = {{3, 7}};
Line(16) = {{4, 8}};

Curve Loop(4) = {{1, 14, -5, -13}};
Plane Surface(4) = {{4}};
Curve Loop(5) = {{2, 15, -6, -14}};
Plane Surface(5) = {{5}};
Curve Loop(6) = {{3, 16, -7, -15}};
Plane Surface(6) = {{6}};
Curve Loop(7) = {{4, 13, -8, -16}};
Plane Surface(7) = {{7}};

// Air box top
Line(17) = {{13, 14}};
Line(18) = {{14, 15}};
Line(19) = {{15, 16}};
Line(20) = {{16, 13}};

Curve Loop(8) = {{17, 18, 19, 20}};
Plane Surface(8) = {{8}};

// Air box sides
Line(21) = {{5, 13}};
Line(22) = {{6, 14}};
Line(23) = {{7, 15}};
Line(24) = {{8, 16}};

Curve Loop(9) = {{5, 22, -17, -21}};
Plane Surface(9) = {{9}};
Curve Loop(10) = {{6, 23, -18, -22}};
Plane Surface(10) = {{10}};
Curve Loop(11) = {{7, 24, -19, -23}};
Plane Surface(11) = {{11}};
Curve Loop(12) = {{8, 21, -20, -24}};
Plane Surface(12) = {{12}};

// Volumes
Surface Loop(1) = {{1, 2, 3, 4, 5, 6, 7}};
Volume(1) = {{1}};  // Substrate

Surface Loop(2) = {{2, 3, 8, 9, 10, 11, 12}};
Volume(2) = {{2}};  // Air

// Physical groups
Physical Surface("Ground", 1) = {{1}};
Physical Surface("Patch", 2) = {{3}};
Physical Surface("ABC", 3) = {{8, 9, 10, 11, 12}};  // Absorbing boundary
Physical Volume("Substrate", 100) = {{1}};  // eps_r = {self.substrate_eps_r}
Physical Volume("Air", 101) = {{2}};

// Mesh settings
Mesh.Algorithm3D = 4;
Mesh.Optimize = 1;
"""
        Path(path).write_text(geo)
        print(f"Generated: {path}")
        return path


@dataclass
class UnitCell:
    """Periodic unit cell for phased array analysis."""
    cell_size_x: float = 0.015  # Cell period in x (m)
    cell_size_y: float = 0.015  # Cell period in y (m)
    element_width: float = 0.008  # Element width
    element_length: float = 0.010  # Element length
    substrate_height: float = 0.00157
    substrate_eps_r: float = 4.4
    air_height: float = 0.015
    mesh_size: float = None

    def __post_init__(self):
        if self.mesh_size is None:
            self.mesh_size = min(self.cell_size_x, self.cell_size_y) / 8

    def generate(self, path: str):
        """Generate Gmsh .geo file for periodic unit cell."""
        cx, cy = self.cell_size_x, self.cell_size_y
        ex, ey = self.element_width, self.element_length
        ox = (cx - ex) / 2
        oy = (cy - ey) / 2
        h = self.substrate_height
        ha = self.air_height

        geo = f"""// Periodic Unit Cell
// Generated by EdgeFEM geometry generator
// Cell: {cx*1000:.2f} x {cy*1000:.2f} mm
// Element: {ex*1000:.2f} x {ey*1000:.2f} mm

lc = {self.mesh_size};
lc_fine = {self.mesh_size/2};

// Ground (z=0)
Point(1) = {{0, 0, 0, lc}};
Point(2) = {{{cx}, 0, 0, lc}};
Point(3) = {{{cx}, {cy}, 0, lc}};
Point(4) = {{0, {cy}, 0, lc}};

Line(1) = {{1, 2}};
Line(2) = {{2, 3}};
Line(3) = {{3, 4}};
Line(4) = {{4, 1}};

Curve Loop(1) = {{1, 2, 3, 4}};
Plane Surface(1) = {{1}};

// Substrate top (z=h)
Point(5) = {{0, 0, {h}, lc}};
Point(6) = {{{cx}, 0, {h}, lc}};
Point(7) = {{{cx}, {cy}, {h}, lc}};
Point(8) = {{0, {cy}, {h}, lc}};

// Element on substrate
Point(9) = {{{ox}, {oy}, {h}, lc_fine}};
Point(10) = {{{ox+ex}, {oy}, {h}, lc_fine}};
Point(11) = {{{ox+ex}, {oy+ey}, {h}, lc_fine}};
Point(12) = {{{ox}, {oy+ey}, {h}, lc_fine}};

// Air top (z=h+ha)
Point(13) = {{0, 0, {h+ha}, lc}};
Point(14) = {{{cx}, 0, {h+ha}, lc}};
Point(15) = {{{cx}, {cy}, {h+ha}, lc}};
Point(16) = {{0, {cy}, {h+ha}, lc}};

// Substrate top edges
Line(5) = {{5, 6}};
Line(6) = {{6, 7}};
Line(7) = {{7, 8}};
Line(8) = {{8, 5}};

// Element edges
Line(9) = {{9, 10}};
Line(10) = {{10, 11}};
Line(11) = {{11, 12}};
Line(12) = {{12, 9}};

// Substrate top surface (with element hole)
Curve Loop(2) = {{5, 6, 7, 8}};
Curve Loop(3) = {{9, 10, 11, 12}};
Plane Surface(2) = {{2, 3}};
Plane Surface(3) = {{3}};  // Element

// Air top
Line(13) = {{13, 14}};
Line(14) = {{14, 15}};
Line(15) = {{15, 16}};
Line(16) = {{16, 13}};

Curve Loop(4) = {{13, 14, 15, 16}};
Plane Surface(4) = {{4}};

// Vertical edges (substrate)
Line(17) = {{1, 5}};
Line(18) = {{2, 6}};
Line(19) = {{3, 7}};
Line(20) = {{4, 8}};

// Vertical edges (air)
Line(21) = {{5, 13}};
Line(22) = {{6, 14}};
Line(23) = {{7, 15}};
Line(24) = {{8, 16}};

// Substrate sides
Curve Loop(5) = {{1, 18, -5, -17}};  // y=0
Plane Surface(5) = {{5}};
Curve Loop(6) = {{2, 19, -6, -18}};  // x=cx
Plane Surface(6) = {{6}};
Curve Loop(7) = {{3, 20, -7, -19}};  // y=cy
Plane Surface(7) = {{7}};
Curve Loop(8) = {{4, 17, -8, -20}};  // x=0
Plane Surface(8) = {{8}};

// Air sides
Curve Loop(9) = {{5, 22, -13, -21}};   // y=0
Plane Surface(9) = {{9}};
Curve Loop(10) = {{6, 23, -14, -22}};  // x=cx
Plane Surface(10) = {{10}};
Curve Loop(11) = {{7, 24, -15, -23}};  // y=cy
Plane Surface(11) = {{11}};
Curve Loop(12) = {{8, 21, -16, -24}};  // x=0
Plane Surface(12) = {{12}};

// Volumes
Surface Loop(1) = {{1, 2, 3, 5, 6, 7, 8}};
Volume(1) = {{1}};

Surface Loop(2) = {{2, 3, 4, 9, 10, 11, 12}};
Volume(2) = {{2}};

// Physical groups
Physical Surface("Ground", 1) = {{1}};
Physical Surface("Element", 2) = {{3}};
Physical Surface("FloquetPort", 3) = {{4}};
Physical Surface("MasterX", 4) = {{8, 12}};  // x=0
Physical Surface("SlaveX", 5) = {{6, 10}};   // x=cx
Physical Surface("MasterY", 6) = {{5, 9}};   // y=0
Physical Surface("SlaveY", 7) = {{7, 11}};   // y=cy
Physical Volume("Substrate", 100) = {{1}};   // eps_r = {self.substrate_eps_r}
Physical Volume("Air", 101) = {{2}};

// Mesh settings
Mesh.Algorithm3D = 4;
Mesh.Optimize = 1;
"""
        Path(path).write_text(geo)
        print(f"Generated: {path}")
        return path


@dataclass
class DipoleAntenna:
    """Half-wave dipole antenna in free space."""
    length: float = 0.15  # Total dipole length (m)
    radius: float = 0.001  # Wire radius (m)
    frequency: float = 1e9  # Design frequency (Hz)
    box_size: float = None  # Simulation box size (auto if None)
    mesh_size: float = None

    def __post_init__(self):
        c0 = 299792458
        wavelength = c0 / self.frequency
        if self.box_size is None:
            self.box_size = wavelength  # 1 wavelength box
        if self.mesh_size is None:
            self.mesh_size = wavelength / 10

    def generate(self, path: str):
        """Generate Gmsh .geo file for dipole antenna."""
        L = self.length
        r = self.radius
        box = self.box_size

        geo = f"""// Dipole Antenna
// Generated by EdgeFEM geometry generator
// Length: {L*1000:.2f} mm, Design freq: {self.frequency/1e9:.3f} GHz

lc = {self.mesh_size};
lc_fine = {self.mesh_size/4};

// Simulation box
Point(1) = {{-{box/2}, -{box/2}, -{box/2}, lc}};
Point(2) = {{ {box/2}, -{box/2}, -{box/2}, lc}};
Point(3) = {{ {box/2},  {box/2}, -{box/2}, lc}};
Point(4) = {{-{box/2},  {box/2}, -{box/2}, lc}};
Point(5) = {{-{box/2}, -{box/2},  {box/2}, lc}};
Point(6) = {{ {box/2}, -{box/2},  {box/2}, lc}};
Point(7) = {{ {box/2},  {box/2},  {box/2}, lc}};
Point(8) = {{-{box/2},  {box/2},  {box/2}, lc}};

// Box edges
Line(1) = {{1, 2}};
Line(2) = {{2, 3}};
Line(3) = {{3, 4}};
Line(4) = {{4, 1}};
Line(5) = {{5, 6}};
Line(6) = {{6, 7}};
Line(7) = {{7, 8}};
Line(8) = {{8, 5}};
Line(9) = {{1, 5}};
Line(10) = {{2, 6}};
Line(11) = {{3, 7}};
Line(12) = {{4, 8}};

// Box faces
Curve Loop(1) = {{1, 2, 3, 4}};
Plane Surface(1) = {{1}};
Curve Loop(2) = {{5, 6, 7, 8}};
Plane Surface(2) = {{2}};
Curve Loop(3) = {{1, 10, -5, -9}};
Plane Surface(3) = {{3}};
Curve Loop(4) = {{2, 11, -6, -10}};
Plane Surface(4) = {{4}};
Curve Loop(5) = {{3, 12, -7, -11}};
Plane Surface(5) = {{5}};
Curve Loop(6) = {{4, 9, -8, -12}};
Plane Surface(6) = {{6}};

// Dipole (simplified as line along z-axis)
Point(9) = {{0, 0, -{L/2}, lc_fine}};
Point(10) = {{0, 0, 0, lc_fine}};  // Feed point
Point(11) = {{0, 0, {L/2}, lc_fine}};

Line(13) = {{9, 10}};  // Lower arm
Line(14) = {{10, 11}}; // Upper arm

// Volume
Surface Loop(1) = {{1, 2, 3, 4, 5, 6}};
Volume(1) = {{1}};

// Embed dipole in volume
Line{{13, 14}} In Volume{{1}};

// Physical groups
Physical Line("Dipole", 1) = {{13, 14}};
Physical Line("FeedGap", 2) = {{10}};  // Node at feed point
Physical Surface("ABC", 3) = {{1, 2, 3, 4, 5, 6}};
Physical Volume("Air", 100) = {{1}};

// Mesh settings
Mesh.Algorithm3D = 4;
Mesh.Optimize = 1;
"""
        Path(path).write_text(geo)
        print(f"Generated: {path}")
        return path


def list_generators():
    """List available geometry generators."""
    print("Available geometry generators:")
    print()
    print("  RectWaveguide  - Rectangular waveguide with wave ports")
    print("  PatchAntenna   - Microstrip patch on dielectric substrate")
    print("  UnitCell       - Periodic unit cell for phased array analysis")
    print("  DipoleAntenna  - Half-wave dipole in free space")
    print()
    print("Usage:")
    print("  from tools.geometry.generators import RectWaveguide")
    print("  wg = RectWaveguide(a=0.02286, b=0.01016, length=0.05)")
    print("  wg.generate('my_waveguide.geo')")


if __name__ == "__main__":
    list_generators()
