import getpass
from contextlib import suppress
from functools import cache
from itertools import product

import gdsfactory as gf
import gfhub
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from gdsfactory.gpdk import PDK
from tqdm.notebook import tqdm

PDK.activate()
np.random.seed(42)  # always generate the same data.
user = getpass.getuser()

client = gfhub.Client()

def iv_resistance(resistance_ohm: float, current_mA: np.ndarray) -> np.ndarray:
    """Calculate voltage from resistance and current (Ohm's law).

    Args:
        resistance_ohm: Resistance in Ohms.
        current_mA: Current in milliamps.

    Returns:
        Voltage in millivolts.
    """
    return resistance_ohm * current_mA  # mV = Ω × mA


# Demo: 30Ω resistor with 0-10 mA sweep
df = pd.DataFrame(
    {
        "current_mA": (i := np.linspace(0, 10, 21)),
        "voltage_mV": iv_resistance(resistance_ohm=30, current_mA=i),
    }
)

plt.plot(df.current_mA, df.voltage_mV)
plt.title("IV Curve (30Ω resistor)")
plt.xlabel("Current (mA)")
plt.ylabel("Voltage (mV)")
plt.grid(True)
plt.show()

wafer_ids = ["wafer1", "wafer2", "wafer3"]
dies = [
    {"x": x, "y": y}
    for y in range(-2, 3)
    for x in range(-2, 3)
    if not (abs(y) == 2 and abs(x) == 2)
]

@cache
def load_gds():
    return gf.import_gds("resistance.gds", skip_new_cells=True)


gds = load_gds()
gds

# Delete existing project files
existing_files = client.query_files(tags=["project:resistance", user])
for file in tqdm(existing_files):
    with suppress(RuntimeError):
        client.delete_file(file["id"])

uploaded = client.add_file("resistance.gds", tags=["project:resistance", user])

sheet_resistance = 100  # Ω/□

for wafer, die in tqdm(list(product(wafer_ids, dies))):
    wafer_variation = 0.20  # ±20% wafer-to-wafer variation
    wafer_factor = 1 + wafer_variation * (2 * np.random.rand() - 1)
    die_id = f"{die['x']},{die['y']}"
    for inst in gds.insts:
        # R = Rs × L / W (sheet resistance formula)
        width_um = inst.cell.info.width
        length_um = inst.cell.info.length
        nominal_resistance = sheet_resistance * length_um / width_um

        device_variation = 0.05  # ±5% device-to-device variation
        device_factor = 1 + device_variation * (2 * np.random.rand() - 1)
        resistance = nominal_resistance * wafer_factor * device_factor

        # Simulate defects: open circuits (5%) and short circuits (3%)
        defect_roll = np.random.rand()
        if defect_roll < 0.05:
            resistance = 1e9  # open circuit
        elif defect_roll < 0.08:
            resistance = 0  # short circuit

        voltage_mV = iv_resistance(
            resistance_ohm=resistance, current_mA=df.current_mA.values
        )
        noise = 0.05 * np.random.rand(len(voltage_mV)) * voltage_mV  # 5% noise
        voltage_mV = voltage_mV + noise

        data = pd.DataFrame(
            {
                "current_mA": df.current_mA.values,
                "voltage_mV": voltage_mV,
            }
        )
        client.add_file(
            data,
            tags=[
                user,
                "project:resistance",
                f"wafer:{wafer}",
                f"die:{die_id}",
                f"cell:{inst.cell.name}",
                f"device:{inst.name}",
                f"length:{inst.cell.info.length}",
                f"width:{inst.cell.info.width}",
            ],
            filename="cutback_device.parquet",
        )

data.to_parquet("last_measurement.parquet")
plt.plot(data["current_mA"], data["voltage_mV"])
plt.xlabel("Current (mA)")
plt.ylabel("Voltage (mV)")
plt.title(f"Last device: {resistance:.1f} Ω")
plt.grid(True)
plt.show()