from flask import Flask, render_template, request, send_file, redirect, url_for
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import io

app = Flask(__name__)

# Global data cache
data_cache = {
    "df1": None,       # Golden data
    "limits": {},
    "cols": [],
    "golden_loaded": False
}


def process_golden_file(golden_file):
    """Load and store Golden Excel data + limits."""
    limits_df1 = pd.read_excel(golden_file, nrows=4)
    df1 = pd.read_excel(golden_file)
    df1 = df1.drop([0, 1, 2, 3])
    df1 = df1.apply(pd.to_numeric, errors="coerce")

    # Extract limits
    limits_df1 = limits_df1.drop([0])
    ignore_cols = ["SITE_NUM", "PART_ID", "PASSFG", "SOFT_BIN", "T_TIME", "TEST_NUM"]
    cols_to_plot = [col for col in limits_df1.columns if "_" in col and col not in ignore_cols]
    limits_df1 = limits_df1.drop(columns=ignore_cols)

    limits = {
        col: {"LL": limits_df1.iloc[0][col], "UL": limits_df1.iloc[1][col]}
        for col in limits_df1.columns
    }

    # Store globally
    data_cache["df1"] = df1
    data_cache["limits"] = limits
    data_cache["cols"] = cols_to_plot
    data_cache["golden_loaded"] = True


def process_test_file(test_file):
    """Load and return the Test Excel data."""
    df2 = pd.read_excel(test_file)
    df2 = df2.drop([0, 1, 2, 3])
    df2 = df2.apply(pd.to_numeric, errors="coerce")
    return df2


def generate_plot(df2, col):
    """Generate plot comparing Golden and Test for a specific column."""
    df1, limits = data_cache["df1"], data_cache["limits"]

    plt.figure(figsize=(6, 4))

    # Golden
    x1 = np.arange(1, len(df1[col]) + 1)
    y1 = pd.to_numeric(df1[col], errors="coerce").values
    plt.plot(x1, y1, marker="o", linestyle="-", color="blue", label="Golden")

    # Test
    if col in df2.columns:
        x2 = np.arange(1, len(df2[col]) + 1)
        y2 = pd.to_numeric(df2[col], errors="coerce").values
        plt.plot(x2, y2, marker="s", linestyle="--", color="red", label="Test")

    # Limits
    if col in limits:
        ll, ul = limits[col]["LL"], limits[col]["UL"]
        plt.axhline(ll, color="green", linestyle="--", linewidth=2, label="LL")
        plt.axhline(ul, color="orange", linestyle="--", linewidth=2, label="UL")

    plt.title(f"{col}")
    plt.xlabel("Part # (sequence)")
    plt.ylabel("Value")
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.legend(fontsize="small")
    plt.tight_layout()

    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight")
    buf.seek(0)
    plt.close()
    return buf


@app.route("/", methods=["GET", "POST"])
def index():
    """Handle golden upload (if not loaded) or test upload (if golden loaded)."""
    if request.method == "POST":
        if not data_cache["golden_loaded"]:
            # Upload Golden
            golden_file = request.files.get("golden_file")
            if not golden_file:
                return render_template("index.html", error="Please upload a Golden file first.")
            try:
                process_golden_file(golden_file)
                return render_template("index.html", message="Golden data loaded successfully! Now upload Test data.")
            except Exception as e:
                return render_template("index.html", error=f"Error loading Golden file: {e}")
        else:
            # Upload Test
            test_file = request.files.get("test_file")
            if not test_file:
                return render_template("index.html", error="Please upload a Test file.")
            try:
                df2 = process_test_file(test_file)
                cols = data_cache["cols"]
                preview_cols = cols[:3] if len(cols) >= 3 else cols
                # Store test temporarily in memory (for this view only)
                data_cache["df2_temp"] = df2
                return render_template("plot.html", cols=cols, preview_cols=preview_cols)
            except Exception as e:
                return render_template("index.html", error=f"Error loading Test file: {e}")

    return render_template("index.html", golden_loaded=data_cache["golden_loaded"])


@app.route("/plot_image/<col>")
def plot_image(col):
    """Generate plot image for selected column."""
    df2 = data_cache.get("df2_temp")
    if df2 is None:
        return "No Test data uploaded yet."
    try:
        buf = generate_plot(df2, col)
        return send_file(buf, mimetype="image/png")
    except Exception as e:
        return f"Error generating plot: {str(e)}"


@app.route("/reset_golden")
def reset_golden():
    """Clear the Golden file from memory."""
    data_cache.update({"df1": None, "limits": {}, "cols": [], "golden_loaded": False})
    return redirect(url_for("index"))


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)