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
import os

app = Flask(__name__)

# Cache
data_cache = {
    "df1": None,
    "limits": {},
    "cols": [],
    "golden_loaded": False,
    "comparison_file": None
}


def process_golden_file(golden_file):
    """Load Golden data and extract 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")

    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
    }

    data_cache.update({
        "df1": df1,
        "limits": limits,
        "cols": cols_to_plot,
        "golden_loaded": True
    })


def process_test_file(test_file):
    """Load Test 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_comparison_excel(df2):
    """Generate comparison Excel (mean, std, min, max for both)."""
    df1 = data_cache["df1"]
    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]
    # common_cols = [ignore_cols = ["SITE_NUM", "PART_ID", "PASSFG", "SOFT_BIN", "T_TIME", "TEST_NUM"]
    common_cols = [col for col in df1.columns if "_" in col and col not in ignore_cols]
    # common_cols = [c for c in df1.columns if c in df2.columns]

    summary = []
    for col in common_cols:
        g_mean, t_mean = df1[col].mean(), df2[col].mean()
        g_std, t_std = df1[col].std(), df2[col].std()
        g_min, t_min = df1[col].min(), df2[col].min()
        g_max, t_max = df1[col].max(), df2[col].max()

        diff = t_mean - g_mean if pd.notna(t_mean) and pd.notna(g_mean) else np.nan
        summary.append([col, g_mean, t_mean, diff, g_std, t_std, g_min, t_min, g_max, t_max])

    comp_df = pd.DataFrame(summary, columns=[
        "Parameter", "Golden_Mean", "Test_Mean", "Mean_Diff",
        "Golden_Std", "Test_Std", "Golden_Min", "Test_Min", "Golden_Max", "Test_Max"
    ])

    path = "comparison_result.xlsx"
    comp_df.to_excel(path, index=False)
    data_cache["comparison_file"] = path


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

    plt.figure(figsize=(6, 4))
    x1 = np.arange(1, len(df1[col]) + 1)
    plt.plot(x1, df1[col], 'o-', label="Golden", color='blue')

    if col in df2.columns:
        x2 = np.arange(1, len(df2[col]) + 1)
        plt.plot(x2, df2[col], 's--', label="Test", color='red')

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

    plt.title(f"{col}")
    plt.xlabel("Part # (sequence)")
    plt.ylabel("Value")
    plt.legend(fontsize='small')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.xticks(np.arange(1, len(df1[col]) + 1))
    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():
    if request.method == "POST":
        # Upload Golden first
        if not data_cache["golden_loaded"]:
            golden_file = request.files.get("golden_file")
            if not golden_file:
                return render_template("index.html", error="Please upload Golden file.")
            try:
                process_golden_file(golden_file)
                return redirect(url_for("index"))
                # return render_template("index.html", message="Golden data loaded successfully!")
            except Exception as e:
                return render_template("index.html", error=f"Error loading Golden file: {e}")

        # Upload Test data next
        else:
            test_file = request.files.get("test_file")
            if not test_file:
                return render_template("index.html", error="Please upload Test data.")
            try:
                df2 = process_test_file(test_file)
                data_cache["df2_temp"] = df2
                generate_comparison_excel(df2)
                return render_template(
                    "plot.html",
                    cols=data_cache["cols"],
                    file_ready=True
                )
            except Exception as e:
                return render_template("index.html", error=f"Error processing Test file: {e}")

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


@app.route("/plot_image/<col>")
def plot_image(col):
    df2 = data_cache.get("df2_temp")
    if df2 is None:
        return "No Test data loaded."
    buf = generate_plot(df2, col)
    return send_file(buf, mimetype="image/png")


@app.route("/download_comparison")
def download_comparison():
    """Download comparison Excel file."""
    path = data_cache.get("comparison_file")
    if path and os.path.exists(path):
        return send_file(path, as_attachment=True)
    return "No comparison file available."


@app.route("/reset_golden")
def reset_golden():
    """Reset golden data."""
    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)