import getpass
import json
from pathlib import Path

import gfhub
import matplotlib.pyplot as plt
import pandas as pd
from gfhub import nodes
from PIL import Image
from scipy import stats
from tqdm.auto import tqdm

user = getpass.getuser()

client = gfhub.Client()

def linear_fit(
    path: Path,
    /,
    *,
    xname: str,
    yname: str,
    slopename: str = "resistance",
    xlabel: str = "",
    ylabel: str = "",
) -> tuple[Path, Path]:
    """Perform linear fit on IV data to extract resistance.

    Args:
        path: Path to parquet file with IV data
        xname: Column name for x-axis (independent variable)
        yname: Column name for y-axis (dependent variable)
        slopename: Name for the extracted slope parameter
        xlabel: Label for x-axis in plot
        ylabel: Label for y-axis in plot

    Returns:
        Tuple of (plot_path, results_path)
    """
    # Load data
    df = pd.read_parquet(path)
    x = df[xname].values
    y = df[yname].values

    # Perform linear fit
    slope, intercept, r_value, p_value, std_err = stats.linregress(x, y)

    # Create plot
    _fig, ax = plt.subplots(figsize=(8, 6))
    ax.scatter(x, y, alpha=0.6, label="Data")
    ax.plot(
        x, slope * x + intercept, "r-", label=f"Fit: y = {slope:.3e}x + {intercept:.3e}"
    )
    xlabel = xlabel or xname
    ylabel = ylabel or yname
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(f"{slopename} = {slope:.3e} (R² = {r_value**2:.4f})")
    ax.legend()
    ax.grid(True)

    # Save plot
    plot_path = path.with_name(path.stem + "_linear_fit.png")
    plt.savefig(plot_path, bbox_inches="tight", dpi=100)
    plt.close()

    # Save results as JSON
    results = {
        slopename: float(slope),
        "intercept": float(intercept),
        "r_squared": float(r_value**2),
        "p_value": float(p_value),
        "std_err": float(std_err),
    }

    results_path = path.with_name(path.stem + "_linear_fit.json")
    results_path.write_text(json.dumps(results, indent=2))

    return plot_path, results_path

func_def = gfhub.Function(
    linear_fit,
    dependencies={
        "pathlib": "from pathlib import Path",
        "json": "import json",
        "matplotlib": "import matplotlib.pyplot as plt",
        "numpy": "import numpy as np",
        "pandas[pyarrow]": "import pandas as pd",
        "scipy": "from scipy import stats",
    },
)

# this should have been created in the previous notebook:
path = Path("last_measurement.parquet").resolve()

result = func_def.eval(path, xname="current_mA", yname="voltage_mV")
print(result)
Image.open(result["output"][0])

client.add_function(func_def)

p = gfhub.Pipeline()

# we can manually trigger the pipeline
p.trigger = nodes.on_manual_trigger()

# or it will auto trigger when a file with these tags gets uploaded.
# note that some of these tags don't have specified parameters values,
# which means the trigger will activate for any of those values.
p.auto_trigger = nodes.on_file_upload(
    tags=[
        ".parquet",
        user,
        "project:resistance",
        "wafer",
        "die",
        "cell",
        "device",
        "length",
        "width",
    ]
)

# the triggers should trigger a file load and a tags load:
p.load_file = nodes.load()
p.load_tags = nodes.load_tags()

# We connect nodes together with the `>>` operator:
p += p.trigger >> p.load_file
p += p.trigger >> p.load_tags
p += p.auto_trigger >> p.load_file
p += p.auto_trigger >> p.load_tags

# after the file is loaded on disk, we'd like to run the analysis function:
p.fit = nodes.function(
    function="linear_fit",
    kwargs={
        "xname": "current_mA",
        "yname": "voltage_mV",
        "slopename": "resistance",
    },
)
p += p.load_file >> p.fit

# the fit function has two outputs: the plot and the json:
p.save_plot = nodes.save()
p.save_json = nodes.save()
# when a node has multiple ports they can be found using their index.
p += p.fit[0] >> p.save_plot[0]
p += p.fit[1] >> p.save_json[0]

# the save nodes also have an input for tags:
p += p.load_tags >> p.save_plot[1]
p += p.load_tags >> p.save_json[1]

pipeline_id = client.add_pipeline(name="device_linear_fit", schema=p)["id"]

client.pipeline_url(pipeline_id)

device_files = client.query_files(
    tags=[
        ".parquet",
        user,
        "project:resistance",
        "wafer",
        "die",
        "cell",
        "device",
        "length",
        "width",
    ]
)
print(f"Found {len(device_files)} device files")

job_ids = []
for device_file in tqdm(device_files):
    triggered = client.trigger_pipeline("device_linear_fit", device_file["id"])
    job_ids.extend(triggered["job_ids"])

print(f"Triggered {len(job_ids)} analysis jobs")

jobs = client.wait_for_jobs(job_ids)

# Query analysis plots
analysis_plots = client.query_files(
    name="*_linear_fit.png", tags=["project:resistance", user]
)

print(f"Found {len(analysis_plots)} analysis plots")

# Display the first plot
if analysis_plots:
    img = Image.open(client.download_file(analysis_plots[0]["id"]))
    display(img.resize((530, 400)))

# Query JSON results
analysis_results = client.query_files(
    name="*_linear_fit.json", tags=["project:resistance", user]
)

print(f"Found {len(analysis_results)} analysis result files")

if analysis_results:
    result_buf = client.download_file(analysis_results[0]["id"])
    result_data = json.load(result_buf)
    print(json.dumps(result_data, indent=2))