import getpass
import json
from pathlib import Path

import gfhub
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from gfhub import nodes
from PIL import Image
from tqdm.notebook import tqdm

client = gfhub.Client()
user = getpass.getuser()
print(f"Running as user: {user}")

entries = client.query_files(
    tags=[
        user,
        "die:-1,-1",       # target a specific die
        "device",
        "project:tutorial_cutback",
        "cell",
        "wafer",
        ".parquet",
    ]
)
print(f"Found {len(entries)} files for die (-1,-1)")

# Download all files for this die
paths = [
    client.download_file(entry["id"], f"file_{i}.parquet")
    for i, entry in enumerate(entries)
]

# Reconstruct tag strings from the entry dict.
# Tags with a parameter_value become "key:value"; simple presence tags become just "key".
tags = [
    [
        (k if not (p := v.get("parameter_value")) else f"{k}:{p}")
        for k, v in entry["tags"].items()
    ]
    for entry in entries
]

print(f"Tags on first file: {tags[0]}")

# Extract component count from tags (stored as "components:816", etc.)
num_comps = [
    int(next(t.replace("components:", "") for t in ts if t.startswith("components:")))
    for ts in tags
]

# Take the peak power across the wavelength axis for each file
powers = [pd.read_parquet(p)["power [dB]"].max() for p in paths]

# Linear fit. polyfit returns [slope, intercept]; both negated because loss reduces power.
component_loss, insertion_loss = [-float(x) for x in np.polyfit(num_comps, powers, deg=1)]

x = np.arange(0, max(num_comps) + 99, 100)
plt.scatter(num_comps, powers, color="C1", label="measured")
plt.plot(x, -component_loss * x - insertion_loss, color="C0", label="fit")
plt.grid(visible=True)
plt.xlim(x.min() - 30, x.max() + 30)
plt.title(f"Die (-1,-1): loss = {component_loss:.4f} dB/component")
plt.xlabel("# components")
plt.ylabel("Power [dB]")
plt.legend()
plt.show()

print(f"Loss per MMI:  {component_loss:.4f} dB")
print(f"Insertion loss at zero components: {insertion_loss:.4f} dB")

def cutback_die_analysis(
    files: list[Path],
    tags: list[list[str]],
    /,
    *,
    output_name: str = "cutback_die_analysis",
) -> tuple[Path, Path]:
    """Run a cutback fit on a group of measurement files from a single die.

    Reads all Parquet files, extracts component counts from tags, fits a line
    through (components, peak power), and saves a plot and a JSON summary.
    """
    dfs = [pd.read_parquet(f) for f in files]

    # Component counts are stored as tags, not inside the files
    num_comps = [
        int(next(t.replace("components:", "") for t in ts if t.startswith("components:")))
        for ts in tags
    ]
    powers = [df["power [dB]"].max() for df in dfs]

    component_loss, insertion_loss = [
        -float(x) for x in np.polyfit(num_comps, powers, deg=1)
    ]

    x = np.arange(0, max(num_comps) + 99, 100)
    plt.scatter(num_comps, powers, color="C1")
    plt.plot(x, -component_loss * x - insertion_loss, color="C0")
    plt.grid(visible=True)
    plt.xlim(x.min() - 30, x.max() + 30)
    plt.title(f"loss = {component_loss:.2e} dB/component")
    plt.xlabel("# components")
    plt.ylabel("Power [dB]")

    path_plot = files[0].parent / f"{output_name}.png"
    plt.savefig(path_plot, bbox_inches="tight")
    plt.close()

    # Parse die coordinates from the first file's tags
    tag_dict = {k.split(":")[0]: k.split(":")[1] for k in tags[0] if ":" in k}
    die_x, die_y = [int(xy) for xy in tag_dict["die"].split(",")]

    output = {
        "die_x": die_x,
        "die_y": die_y,
        "component_loss": None if not np.isfinite(component_loss) else component_loss,
        "insertion_loss": None if not np.isfinite(insertion_loss) else insertion_loss,
    }
    path_json = files[0].parent / f"{output_name}.json"
    path_json.write_text(json.dumps(output))

    return path_plot, path_json

func_def = gfhub.Function(
    cutback_die_analysis,
    dependencies={
        "json": "import json",
        "numpy": "import numpy as np",
        "pandas[pyarrow]": "import pandas as pd",
        "matplotlib": "import matplotlib.pyplot as plt",
    },
)

result = func_def.eval(paths, tags)
print(result)  # {"success": True, "output": ["/path/to/plot.png", "/path/to/analysis.json"]}

Image.open(result["output"][0])

client.add_function(func_def)
print("cutback_die_analysis uploaded.")

def find_common_tags(tags: list[list[str]], /) -> list[str]:
    """Return only the tags that are identical across all input files.

    Tags that vary between files are dropped. Extension tags (keys starting
    with '.') are always excluded regardless of whether they match.
    """
    # Build a dict mapping each tag key to the set of values seen across all files
    common: dict[str, set] = {}
    for _tags in tags:
        for t in _tags:
            if ":" in t:
                key, value = t.split(":", 1)
            else:
                key, value = t, ""
            common.setdefault(key, set()).add(value)

    # Keep only keys where all files agreed on the same value
    agreed = {k: next(iter(v)) for k, v in common.items() if len(v) == 1}

    # Reconstruct as tag strings, skipping extension keys
    return [
        k if not v else f"{k}:{v}"
        for k, v in agreed.items()
        if not k.startswith(".")
    ]

# Verify: die, wafer, and project should be shared; components and cell should not appear
common = find_common_tags(tags)
print(common)

client.add_function(find_common_tags)
print("find_common_tags uploaded.")

p = gfhub.Pipeline()

# Manual trigger only: the fit requires a full group of files, not a single upload
p.trigger = nodes.on_manual_trigger()

# Download all files in the group and their tags
p.load_file = nodes.load()
p.load_tags = nodes.load_tags()
p += p.trigger >> p.load_file
p += p.trigger >> p.load_tags

# Run the cutback fit: files go to input 0, tags go to input 1
p.cutback_die_analysis = nodes.function(function="cutback_die_analysis")
p += p.load_file >> p.cutback_die_analysis[0]
p += p.load_tags >> p.cutback_die_analysis[1]

# Find the tags shared by all input files, to attach to both outputs
p.common_tags = nodes.function(function="find_common_tags")
p += p.load_tags >> p.common_tags

# Save the plot (output 0) and the JSON (output 1), both with the shared tags
p.save_plot = nodes.save()
p += p.cutback_die_analysis[0] >> p.save_plot[0]
p += p.common_tags >> p.save_plot[1]

p.save_json = nodes.save()
p += p.cutback_die_analysis[1] >> p.save_json[0]
p += p.common_tags >> p.save_json[1]

confirmation = client.add_pipeline(name="cutback_die_analysis", schema=p)
print(f"Pipeline ready: {client.pipeline_url(confirmation['id'])}")

entries = client.query_files(
    tags=[user, "die", "device", "project:tutorial_cutback", "cell", "wafer", ".parquet"]
).groupby("die")

job_ids = []
for die_tag, group in tqdm(entries.items()):
    input_ids = [f["id"] for f in group]
    triggered = client.trigger_pipeline("cutback_die_analysis", input_ids)
    job_ids.extend(triggered["job_ids"])

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

jobs = client.wait_for_jobs(job_ids)
print(f"All jobs complete. Statuses: {[j['status'] for j in jobs]}")