main.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import torch
import torch.nn.functional as F
import cv2
import numpy as np
from PIL import Image
import io
import base64
from transformers import AutoImageProcessor, AutoModel, CLIPProcessor, CLIPModel
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from scipy.optimize import linear_sum_assignment
import lpips
import torchvision.transforms as transforms

# ==============================================================================
# 1. Initialize FastAPI & CORS
# ==============================================================================
app = FastAPI(title="Copyright Diagnostic API - 3 Pillar XAI")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==============================================================================
# 2. Global Initialization & Memory Management
# ==============================================================================
device = "cuda" if torch.cuda.is_available() else "cpu"

DINO_MODEL_ID = "facebook/dinov2-large"
CLIP_MODEL_ID = "openai/clip-vit-large-patch14"
PATCH_SIZE = 14

print(f"Loading DINOv2 ({DINO_MODEL_ID})...")
dino_processor = AutoImageProcessor.from_pretrained(DINO_MODEL_ID)
dino_model = AutoModel.from_pretrained(DINO_MODEL_ID).to(device)
dino_model.eval()

print(f"Loading CLIP ({CLIP_MODEL_ID})...")
clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
clip_model = CLIPModel.from_pretrained(CLIP_MODEL_ID).to(device)
clip_model.eval()

print("Loading LPIPS (AlexNet)...")
loss_fn_alex = lpips.LPIPS(net='alex').to(device)
loss_fn_alex.eval()

# LPIPS requires images normalized between [-1, 1]
lpips_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

# ==============================================================================
# 3. Helpers & Base64 Converters
# ==============================================================================

def image_to_base64(img: Image.Image) -> str:
    buffered = io.BytesIO()
    img.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def fig_to_base64(fig) -> str:
    buf = io.BytesIO()
    fig.savefig(buf, format="jpg", bbox_inches='tight', pad_inches=0.1, dpi=100)
    plt.close(fig)
    return base64.b64encode(buf.getvalue()).decode("utf-8")

def to_edge_map(image: Image.Image) -> Image.Image:
    """Strips style/texture, leaving structural contours for Pillar 1."""
    img_array = np.array(image.convert("L"))
    median = np.median(img_array)
    low = int(max(0, 0.5 * median))
    high = int(min(255, 1.3 * median))
    edges = cv2.Canny(img_array, low, high)
    kernel = np.ones((2, 2), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)
    return Image.fromarray(edges).convert("RGB")

def patch_idx_to_xy(idx, grid_w, patch_size):
    row = idx // grid_w
    col = idx % grid_w
    return col * patch_size + patch_size / 2, row * patch_size + patch_size / 2

# ==============================================================================
# 4. AI Pipeline Functions (Mapped to the 3 Pillars)
# ==============================================================================

def extract_dino_features(image: Image.Image, preprocess="color"):
    if preprocess == "grayscale":
        img = image.convert("L").convert("RGB")
    elif preprocess == "edges":
        img = to_edge_map(image)
    else:
        img = image

    inputs = dino_processor(images=img, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = dino_model(**inputs)

    cls_token = outputs.last_hidden_state[:, 0, :]
    patch_tokens = outputs.last_hidden_state[:, 1:, :]
    n_patches = patch_tokens.shape[1]
    grid_size = int(n_patches ** 0.5)
    return cls_token, patch_tokens, grid_size, grid_size

def extract_clip_similarity(image_a: Image.Image, image_b: Image.Image):
    inputs = clip_processor(images=[image_a, image_b], return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = clip_model.get_image_features(**inputs)
        
    if hasattr(outputs, 'image_embeds'):
        features = outputs.image_embeds
    elif isinstance(outputs, torch.Tensor):
        features = outputs
    else:
        features = outputs.pooler_output
        
    features = F.normalize(features, dim=-1)
    return round((features[0] @ features[1]).item(), 4)

def compute_patch_matches(patches_a, patches_b):
    pa = F.normalize(patches_a.squeeze(0), dim=-1)
    pb = F.normalize(patches_b.squeeze(0), dim=-1)
    sim_matrix = pa @ pb.T
    a_to_b_scores = sim_matrix.max(dim=1).values
    b_to_a_scores = sim_matrix.max(dim=0).values
    return sim_matrix, a_to_b_scores, b_to_a_scores

def nms_matches(matches, grid_w, patch_size, radius=2.0):
    if not matches: return []
    matches = sorted(matches, key=lambda m: m[2], reverse=True)
    kept = []
    pixel_radius = radius * patch_size

    for idx_a, idx_b, score in matches:
        xa, ya = patch_idx_to_xy(idx_a, grid_w, patch_size)
        xb, yb = patch_idx_to_xy(idx_b, grid_w, patch_size)
        dominated = False
        for ka, kb, ks in kept:
            kxa, kya = patch_idx_to_xy(ka, grid_w, patch_size)
            kxb, kyb = patch_idx_to_xy(kb, grid_w, patch_size)
            if (abs(xa - kxa) < pixel_radius and abs(ya - kya) < pixel_radius) or \
               (abs(xb - kxb) < pixel_radius and abs(yb - kyb) < pixel_radius):
                dominated = True
                break
        if not dominated:
            kept.append((idx_a, idx_b, score))
    return kept

def make_correspondence_figure(image_a, image_b, patches_a, patches_b, grid_h, grid_w, max_matches=20, score_thresh=0.5):
    pa = patches_a.squeeze(0).cpu()
    pb = patches_b.squeeze(0).cpu()

    pa_norm = F.normalize(pa, dim=-1)
    pb_norm = F.normalize(pb, dim=-1)

    sim = (pa_norm @ pb_norm.T).numpy()
    cost = 1.0 - sim

    row_ind, col_ind = linear_sum_assignment(cost)

    raw_matches = []
    for r, c in zip(row_ind, col_ind):
        score = sim[r, c]
        if score >= score_thresh:
            raw_matches.append((r, c, float(score)))

    # 1. Calculate ALL valid matches that pass the threshold
    all_valid_matches = nms_matches(raw_matches, grid_w, PATCH_SIZE, radius=2.0)
    total_match_count = len(all_valid_matches)

    # 2. Slice the list to only draw the top N matches to prevent visual clutter
    vis_matches = all_valid_matches[:max_matches]

    img_w, img_h = grid_w * PATCH_SIZE, grid_h * PATCH_SIZE
    img_a_resized = image_a.resize((img_w, img_h))
    img_b_resized = image_b.resize((img_w, img_h))

    gap = 30
    canvas_w = img_w * 2 + gap
    fig, ax = plt.subplots(1, 1, figsize=(14, 6))
    fig.patch.set_facecolor('#0f172a') 

    canvas = Image.new("RGB", (canvas_w, img_h), (15, 23, 42)) 
    canvas.paste(img_a_resized, (0, 0))
    canvas.paste(img_b_resized, (img_w + gap, 0))
    ax.imshow(canvas)

    cmap = plt.cm.get_cmap("spring", max(len(vis_matches), 1))

    # 3. Only iterate over the sliced vis_matches for drawing
    for i, (idx_a, idx_b, score) in enumerate(vis_matches):
        xa, ya = patch_idx_to_xy(idx_a, grid_w, PATCH_SIZE)
        xb, yb = patch_idx_to_xy(idx_b, grid_w, PATCH_SIZE)
        xb_canvas = xb + img_w + gap
        color = cmap(i % 20)

        tl_xa = xa - PATCH_SIZE / 2
        tl_ya = ya - PATCH_SIZE / 2
        tl_xb = xb_canvas - PATCH_SIZE / 2
        tl_yb = yb - PATCH_SIZE / 2

        rect_a = plt.Rectangle((tl_xa, tl_ya), PATCH_SIZE, PATCH_SIZE,
                               linewidth=1.5, edgecolor='#ef4444', facecolor='none', alpha=0.9, zorder=4)
        ax.add_patch(rect_a)

        rect_b = plt.Rectangle((tl_xb, tl_yb), PATCH_SIZE, PATCH_SIZE,
                               linewidth=1.5, edgecolor='#ef4444', facecolor='none', alpha=0.9, zorder=4)
        ax.add_patch(rect_b)

        ax.plot([xa, xb_canvas], [ya, yb], color=color, linewidth=2, alpha=0.8)
        ax.scatter([xa, xb_canvas], [ya, yb], color=color, s=50, zorder=5, edgecolors="white", linewidths=0.5)

        mx = (xa + xb_canvas) / 2
        my = (ya + yb) / 2
        
        bbox_props = dict(boxstyle="round,pad=0.25", fc="#1e293b", ec=color, alpha=0.95, lw=1)
        ax.text(mx, my, f"{score:.2f}", fontsize=8, color="white", fontweight="bold",
                ha="center", va="center", zorder=10, bbox=bbox_props)

    ax.axis("off")
    fig.tight_layout(pad=0)
    
    # 4. Return the TOTAL count, not the sliced count
    return fig, total_match_count

def make_combined_figure(image_a, image_b, scores_a, scores_b, grid_h, grid_w):
    heatmap_a = scores_a.reshape(grid_h, grid_w).cpu().numpy()
    heatmap_b = scores_b.reshape(grid_h, grid_w).cpu().numpy()

    fig = plt.figure(figsize=(12, 5.5))
    fig.patch.set_facecolor('#0f172a') # Dark mode mapping
    gs = gridspec.GridSpec(1, 2, wspace=0.05)

    ax0 = fig.add_subplot(gs[0])
    ax0.imshow(image_a.resize((grid_w * PATCH_SIZE, grid_h * PATCH_SIZE)))
    ax0.imshow(heatmap_a, cmap="inferno", alpha=0.6, interpolation="bilinear", extent=(0, grid_w * PATCH_SIZE, grid_h * PATCH_SIZE, 0))
    ax0.axis("off")

    ax1 = fig.add_subplot(gs[1])
    ax1.imshow(image_b.resize((grid_w * PATCH_SIZE, grid_h * PATCH_SIZE)))
    ax1.imshow(heatmap_b, cmap="inferno", alpha=0.6, interpolation="bilinear", extent=(0, grid_w * PATCH_SIZE, grid_h * PATCH_SIZE, 0))
    ax1.axis("off")

    fig.tight_layout(pad=0)
    return fig

# ==============================================================================
# 5. Primary Analysis Endpoint
# ==============================================================================

@app.post("/analyze")
async def analyze_artworks(file_a: UploadFile = File(...), file_b: UploadFile = File(...)):
    try:
        img_a = Image.open(io.BytesIO(await file_a.read())).convert("RGB")
        img_b = Image.open(io.BytesIO(await file_b.read())).convert("RGB")
        
        # --- PILLAR 1: Idea-Expression Dichotomy ---
        semantic_idea_score = extract_clip_similarity(img_a, img_b)
        
        cls_e_a, patches_e_a, gh, gw = extract_dino_features(img_a, "edges")
        cls_e_b, patches_e_b, _, _ = extract_dino_features(img_b, "edges")
        structural_expression_score = round(F.cosine_similarity(cls_e_a, cls_e_b).item(), 4)
        
        edge_a_b64 = image_to_base64(to_edge_map(img_a))
        edge_b_b64 = image_to_base64(to_edge_map(img_b))

        # --- PILLAR 2: Fragmented Literal Similarity (RESTored BEST-OF FUSION) ---
        cls_c_a, patches_c_a, _, _ = extract_dino_features(img_a, "color")
        cls_c_b, patches_c_b, _, _ = extract_dino_features(img_b, "color")
        
        cls_g_a, patches_g_a, _, _ = extract_dino_features(img_a, "grayscale")
        cls_g_b, patches_g_b, _, _ = extract_dino_features(img_b, "grayscale")
        
        cls_e_a, patches_e_a, _, _ = extract_dino_features(img_a, "edges")
        cls_e_b, patches_e_b, _, _ = extract_dino_features(img_b, "edges")
        
        _, a2b_c, b2a_c = compute_patch_matches(patches_c_a, patches_c_b)
        _, a2b_g, b2a_g = compute_patch_matches(patches_g_a, patches_g_b)
        _, a2b_e, b2a_e = compute_patch_matches(patches_e_a, patches_e_b)

        a2b_best = torch.max(torch.max(a2b_c, a2b_g), a2b_e)
        b2a_best = torch.max(torch.max(b2a_c, b2a_g), b2a_e)
        
        corr_thresh = (a2b_best.mean() + 0.5 * a2b_best.std()).item()
        corr_thresh = min(max(corr_thresh, 0.4), 0.75)
        
        mode_scores = {
            "color": a2b_c.mean().item(),
            "grayscale": a2b_g.mean().item(),
            "edges": a2b_e.mean().item()
        }
        best_mode = max(mode_scores, key=mode_scores.get)
        mode_patches = {
            "color": (patches_c_a, patches_c_b),
            "grayscale": (patches_g_a, patches_g_b),
            "edges": (patches_e_a, patches_e_b),
        }
        corr_pa, corr_pb = mode_patches[best_mode]

        # Call the figure generator exactly ONCE. It now returns the total count.
        corr_fig, total_match_count = make_correspondence_figure(
            img_a, img_b, corr_pa, corr_pb, gh, gw, score_thresh=corr_thresh
        )
        correspondence_map_b64 = fig_to_base64(corr_fig)
        
        n_patches_a = a2b_best.shape[0]
        
        # 1. Calculate the raw statistical threshold
        raw_adaptive_thresh = (a2b_best.mean() + a2b_best.std()).item()
        
        # 2. Clamp it: never lower than 0.5 (noise), never higher than 0.85 (identical images)
        smart_thresh = min(max(raw_adaptive_thresh, 0.5), 0.85)
        
        # 3. Count patches above this dynamic, safe threshold
        high_a = (a2b_best > smart_thresh).sum().item()
        pct_copied = round((high_a / n_patches_a) * 100, 1)

        # --- PILLAR 3: Substantial Similarity ---
        heatmap_fig = make_combined_figure(img_a, img_b, a2b_best, b2a_best, gh, gw)
        heatmap_b64 = fig_to_base64(heatmap_fig)
        
        t_a = lpips_transform(img_a).unsqueeze(0).to(device)
        t_b = lpips_transform(img_b).unsqueeze(0).to(device)
        with torch.no_grad():
            lpips_distance = round(loss_fn_alex(t_a, t_b).item(), 4)

        return {
            "status": "success",
            "pillar_1_idea_expression": {
                "semantic_idea_score": semantic_idea_score,
                "structural_expression_score": structural_expression_score,
                "edge_map_a_b64": f"data:image/jpeg;base64,{edge_a_b64}",
                "edge_map_b_b64": f"data:image/jpeg;base64,{edge_b_b64}"
            },
            "pillar_2_fragmented_literal": {
                "patch_match_count": total_match_count,
                "percentage_copied": pct_copied,
                "correspondence_map_b64": f"data:image/jpeg;base64,{correspondence_map_b64}"
            },
            "pillar_3_substantial_similarity": {
                "perceptual_distance_lpips": lpips_distance,
                "quantitative_heatmap_b64": f"data:image/jpeg;base64,{heatmap_b64}"
            }
        }
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def read_root():
    return {"message": "3-Pillar Legal Diagnostic API is running."}