import os, cv2, re, base64
import numpy as np
import pandas as pd
import gradio as gr
from roboflow import Roboflow
from openai import OpenAI
from openpyxl import load_workbook

# ================= CONFIG =================
ROBOFLOW_API_KEY = "uP19IAi98TqwLvHmNB8V"
ROBOFLOW_PROJECT = "braker3"
ROBOFLOW_VERSION = 6
CONF_THRESHOLD = 0.35
IOU_THRESHOLD = 0.4

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
rf = Roboflow(api_key=ROBOFLOW_API_KEY)
model = rf.workspace().project(ROBOFLOW_PROJECT).version(ROBOFLOW_VERSION).model

# ================= CONSTANTS =================
SPEC_JP = {
    "Manufacture Name": "メーカー",
    "Circuit Name": "回路番号",
    "Load Name": "負荷名称",
    "Breaking Capacity": "遮断容量",
    "AF": "フレーム(AF)",
    "AT": "トリップ(AT)"
}

MANUFACTURER_MAP = {
    "MITSUBISHI ELECTRIC": "三菱電機",
    "SIEMENS": "SIEMENS",
    "SCHNEIDER ELECTRIC": "SCHNEIDER ELECTRIC",
    "ABB": "ABB",
    "LS ELECTRIC": "LS ELECTRIC"
}

VALID_BREAKING_CAPACITY = {"6","10","15","25","36","50","65","85"}
DEFAULT_BREAKING_CAPACITY = "85"   # ← your dataset truth

# ================= IMAGE =================
def prepare_for_roboflow(img, max_side=1024):
    h,w = img.shape[:2]
    scale = min(max_side/max(h,w),1.0)
    return cv2.resize(img,(int(w*scale),int(h*scale))) if scale<1 else img

def crop(img,x1,y1,x2,y2,pad=20):
    h,w = img.shape[:2]
    return img[max(0,y1-pad):min(h,y2+pad), max(0,x1-pad):min(w,x2+pad)]

def enhance(img):
    g = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    g = cv2.resize(g,None,fx=3,fy=3,interpolation=cv2.INTER_CUBIC)
    clahe = cv2.createCLAHE(2.0,(8,8))
    return cv2.cvtColor(clahe.apply(g), cv2.COLOR_GRAY2BGR)

def enhance_breaking_capacity(img):
    g = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    g = cv2.resize(g,None,fx=4,fy=4,interpolation=cv2.INTER_CUBIC)
    g = cv2.adaptiveThreshold(
        g,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,31,2
    )
    return cv2.cvtColor(g, cv2.COLOR_GRAY2BGR)

def img_to_base64(img):
    _,buf = cv2.imencode(".jpg",img)
    return base64.b64encode(buf).decode()

# ================= TEXT =================
def remove_spaces_only(text):
    return re.sub(r"\s+", "", str(text)) if text else ""

def extract_digits(text):
    nums = re.findall(r"\d+",str(text))
    return nums[-1] if nums else ""

def clean_manufacturer(text):
    t=text.upper()
    for k in MANUFACTURER_MAP:
        if k in t:
            return k
    return ""

def normalize_breaking_capacity(text):
    nums = re.findall(r"\d+",str(text))
    for n in nums:
        if n in VALID_BREAKING_CAPACITY:
            return n
    return DEFAULT_BREAKING_CAPACITY

# ================= GPT OCR =================
def gpt_ocr(label,img):
    if label == "Breaking Capacity":
        img = enhance_breaking_capacity(img)
    else:
        img = enhance(img)

    b64 = img_to_base64(img)

    rules={
        "Manufacture Name":"Return ONLY manufacturer name in English.",
        "Circuit Name":"Return EXACT text.",
        "Load Name":"Return EXACT text.",
        "AF":"Return ONLY number.",
        "AT":"Return ONLY number.",
        "Breaking Capacity":"Return ONLY kA number."
    }

    r = client.chat.completions.create(
        model="gpt-5.2",
        messages=[
            {"role":"system","content":"Strict OCR engine"},
            {"role":"user","content":[
                {"type":"text","text":rules[label]},
                {"type":"image_url","image_url":{"url":f"data:image/jpeg;base64,{b64}"}}
            ]}
        ],
        temperature=0
    )

    raw = r.choices[0].message.content.strip()

    if label == "Manufacture Name":
        return clean_manufacturer(raw)
    if label in ["Circuit Name","Load Name"]:
        return remove_spaces_only(raw)
    if label in ["AF","AT"]:
        return extract_digits(raw)
    if label == "Breaking Capacity":
        return normalize_breaking_capacity(raw)

    return raw

# ================= EXCEL =================
def normalize_header(s): 
    return str(s).replace("\n","").replace(" ","")

def find_column(df,keys):
    for c in df.columns:
        for k in keys:
            if k in normalize_header(c):
                return c
    return None

def verify_excel(excel,det):
    wb=load_workbook(excel,data_only=True)
    ws=wb["MCB"]

    raw=pd.DataFrame([list(r) for r in ws.iter_rows(values_only=True)])
    raw.dropna(how="all",inplace=True)

    hdr=None
    for i in range(len(raw)):
        if "回路" in "".join(map(str,raw.iloc[i].values)):
            hdr=i; break

    if hdr is None:
        return pd.DataFrame([["回路番号","", "NO","ヘッダー不明"]],
            columns=["仕様","検出値","Excelに存在？","備考"])

    df=raw.iloc[hdr+1:].copy()
    df.columns=raw.iloc[hdr]
    df.dropna(how="all",inplace=True)

    ccol=find_column(df,["回路番号","回路"])
    target=None
    for _,r in df.iterrows():
        if remove_spaces_only(r[ccol])==det.get("Circuit Name",""):
            target=r; break

    if target is None:
        return pd.DataFrame([["回路番号",det.get("Circuit Name",""),"NO","Excelに存在しない"]],
            columns=["仕様","検出値","Excelに存在？","備考"])

    rows=[]
    for k,jp in SPEC_JP.items():
        detv=det.get(k,"")
        col=find_column(df,[jp.replace("(","").replace(")",""),jp[:2]])
        excelv=str(target[col]) if col else ""

        if k in ["Circuit Name","Load Name"]:
            ok=remove_spaces_only(detv)==remove_spaces_only(excelv)
        elif k=="Manufacture Name":
            ok=MANUFACTURER_MAP.get(detv,detv)==excelv
        else:
            ok=detv==excelv

        rows.append([jp,detv,"YES" if ok else "NO","" if ok else f"Excel値: {excelv}"])

    return pd.DataFrame(rows,columns=["仕様","検出値","Excelに存在？","備考"])

# ================= PIPELINE =================
def run_pipeline(image,excel):
    if image is None:
        return None,pd.DataFrame(),pd.DataFrame(),None

    img=prepare_for_roboflow(image)
    preds=model.predict(
        img,
        confidence=int(CONF_THRESHOLD*100),
        overlap=int(IOU_THRESHOLD*100)
    ).json()["predictions"]

    best={}
    vis=img.copy()

    for p in preds:
        lab=p["class"]
        x,y,w,h=map(int,[p["x"],p["y"],p["width"],p["height"]])
        x1,y1,x2,y2=x-w//2,y-h//2,x+w//2,y+h//2
        cv2.rectangle(vis,(x1,y1),(x2,y2),(0,255,0),2)
        c=crop(img,x1,y1,x2,y2)
        if lab not in best or p["confidence"]>best[lab][0]:
            best[lab]=(p["confidence"],c)

    det={}
    rows=[]
    for lab,(_,c) in best.items():
        v=gpt_ocr(lab,c)
        if v:
            det[lab]=v
            rows.append([lab,v])

    return vis, pd.DataFrame(rows,columns=["Field","Extracted Text"]), verify_excel(excel,det), "verification_result.xlsx"

# ================= UI =================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ⚡ Breaker Panel OCR & Verification")

    with gr.Row():
        img_in=gr.Image(type="numpy",label="Upload Image")
        img_out=gr.Image(label="Detected Image")

    excel_in=gr.File(label="Upload Excel (MCB)",file_types=[".xlsx",".xlsm"])
    btn=gr.Button("Run Verification",variant="primary")

    t1=gr.Dataframe(label="OCR Output")
    t2=gr.Dataframe(label="Verification Result")
    f=gr.File(label="Download Result")

    btn.click(run_pipeline,[img_in,excel_in],[img_out,t1,t2,f])
    demo.launch()