#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from sklearn.linear_model import Ridge

# =========================
# VECTOR EXPORT SETTINGS
# =========================
mpl.rcParams["pdf.fonttype"] = 42
mpl.rcParams["ps.fonttype"] = 42

# =========================
# PATHS
# =========================
BASE = os.path.dirname(os.path.abspath(__file__))

anno_csv = os.path.join(BASE, "ROI_ANNOTATION_SF9", "roi_labels.csv")
img_folder = os.path.join(BASE, "images")
input_folder = os.path.join(BASE, "SF9rename")

out = os.path.join(BASE, "SF9_PUBLICATION_FINAL")
qc_out = os.path.join(out, "QC")

os.makedirs(out, exist_ok=True)
os.makedirs(qc_out, exist_ok=True)

ROI = 229

# =========================
# SAFE ROI
# =========================
def crop_roi_safe(img, cx, cy, size):

    h, w = img.shape
    half = size // 2

    x1 = int(cx - half)
    x2 = int(cx + half)
    y1 = int(cy - half)
    y2 = int(cy + half)

    pad_l = max(0, -x1)
    pad_t = max(0, -y1)
    pad_r = max(0, x2 - w)
    pad_b = max(0, y2 - h)

    x1 = max(0, x1)
    y1 = max(0, y1)
    x2 = min(w, x2)
    y2 = min(h, y2)

    roi = img[y1:y2, x1:x2]

    roi = cv2.copyMakeBorder(
        roi,
        pad_t, pad_b,
        pad_l, pad_r,
        borderType=cv2.BORDER_REFLECT
    )

    roi = cv2.resize(roi, (size, size))

    return roi

# =========================
# NORMALIZE (no log)
# =========================
def normalize_roi(roi):

    roi = roi.astype(np.float32)
    p1, p99 = np.percentile(roi, (1, 99))
    return np.clip(roi, p1, p99)

# =========================
# FEATURE ENGINEERING
# =========================
def extract_features(roi):

    thr = np.mean(roi) + 0.3*np.std(roi)
    binary = roi > thr

    area = np.sum(binary)
    ratio = area / roi.size
    num_labels, _ = cv2.connectedComponents(binary.astype(np.uint8))
    intensity = np.sum(roi)

    avg_blob = area / (num_labels + 1e-6)

    return np.array([
        area / 1000.0,
        ratio,
        num_labels,
        intensity / 1e6,
        avg_blob / 100.0
    ], dtype=np.float32)

# =========================
# LOAD DATA
# =========================
df = pd.read_csv(anno_csv)

print("\n===== SF9 TRAIN =====")
print("ROI samples:", len(df))

X, y = [], []

for _, r in df.iterrows():

    img = cv2.imread(os.path.join(img_folder, r["image"]), 0)
    if img is None:
        continue

    roi = crop_roi_safe(img, r["x"], r["y"], ROI)
    roi = normalize_roi(roi)

    X.append(extract_features(roi))
    y.append(float(r["count"]))

X = np.array(X)
y = np.array(y)

print("Valid:", len(X))

# =========================
# TRAIN MODEL
# =========================
model = Ridge(alpha=1.0)
model.fit(X, y)

pred = model.predict(X)

r2 = model.score(X, y)

print("\nR2 =", r2)

# =========================
# METRICS
# =========================
diff = pred - y

pd.DataFrame({
    "R2":[r2],
    "mean_error":[np.mean(diff)],
    "std_error":[np.std(diff)]
}).to_csv(os.path.join(out, "metrics_summary.csv"), index=False)

# =========================
# FIGURE 1 - SCATTER (PDF)
# =========================
plt.figure()
plt.scatter(y, pred, s=18)

plt.plot([y.min(), y.max()],
         [y.min(), y.max()],
         'r--')

plt.xlabel("Manual count")
plt.ylabel("Predicted count")
plt.title(f"SF9 regression (R2={r2:.3f})")

plt.tight_layout()
plt.savefig(os.path.join(out, "SF9_scatter.pdf"))
plt.close()

# =========================
# FIGURE 2 - BLAND ALTMAN
# =========================
mean = (pred + y) / 2

plt.figure()
plt.scatter(mean, diff, s=18)

plt.axhline(np.mean(diff), linestyle='--')
plt.axhline(np.mean(diff)+1.96*np.std(diff), linestyle='--')
plt.axhline(np.mean(diff)-1.96*np.std(diff), linestyle='--')

plt.xlabel("Mean of measurements")
plt.ylabel("Difference (Pred - Manual)")
plt.title("Bland–Altman")

plt.tight_layout()
plt.savefig(os.path.join(out, "SF9_BlandAltman.pdf"))
plt.close()

# =========================
# FIGURE 3 - RESIDUAL
# =========================
plt.figure()
plt.hist(diff, bins=12)

plt.xlabel("Error")
plt.ylabel("Frequency")
plt.title("Residual distribution")

plt.tight_layout()
plt.savefig(os.path.join(out, "SF9_residual.pdf"))
plt.close()

# =========================
# ROI QC INFERENCE
# =========================
roi_log = []
summary = []

for fn in os.listdir(input_folder):

    if not fn.endswith(".jpg"):
        continue

    img = cv2.imread(os.path.join(input_folder, fn), 0)
    if img is None:
        continue

    h, w = img.shape

    preds = []
    roi_id = 0

    fig, ax = plt.subplots(figsize=(6,6))
    ax.imshow(img, cmap='gray')

    step = ROI // 2

    for y in range(ROI, h-ROI, step):
        for x in range(ROI, w-ROI, step):

            roi = crop_roi_safe(img, x, y, ROI)
            roi = normalize_roi(roi)

            p = model.predict([extract_features(roi)])[0]
            preds.append(p)

            roi_log.append([fn, roi_id, x, y, p])

            ax.add_patch(plt.Rectangle(
                (x-ROI, y-ROI), ROI, ROI,
                fill=False, edgecolor='red', linewidth=0.6
            ))

            ax.text(x-ROI, y-ROI, f"{p:.1f}", color='yellow', fontsize=5)

            roi_id += 1

    summary.append([fn, np.mean(preds), np.std(preds)])

    ax.set_title(fn)
    ax.axis("off")

    plt.savefig(os.path.join(qc_out, fn + "_QC.png"), dpi=300)
    plt.close()

# =========================
# SAVE TABLES
# =========================
pd.DataFrame(summary, columns=["image","mean","std"]).to_csv(
    os.path.join(out, "summary.csv"), index=False)

pd.DataFrame(roi_log, columns=["image","roi_id","x","y","pred"]).to_csv(
    os.path.join(out, "roi_details.csv"), index=False)

print("\n✅ PUBLICATION PACKAGE COMPLETE")