backend/model_service.py

from __future__ import annotations

import base64
import mimetypes
import io
import json
import os
import sys
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, ImageStat
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from torchvision import models, transforms

ROOT = Path(__file__).resolve().parent.parent
REPO_DIR = ROOT / 'repo_inspect'
INTEGRATION_ASSETS_DIR = ROOT / 'integration_assets'
VIDEO_BUNDLE_DIR = ROOT / 'video_bundle'
VIDEO_CODE_DIR = VIDEO_BUNDLE_DIR / 'Video'
for path in [ROOT, VIDEO_CODE_DIR, VIDEO_BUNDLE_DIR, REPO_DIR, INTEGRATION_ASSETS_DIR]:
    path_str = str(path)
    if path_str in sys.path:
        sys.path.remove(path_str)
    sys.path.insert(0, path_str)

from ethical_assessment import EthicalAssessment, format_ethical_report, get_simple_status  # noqa: E402
from video_model import ResNetLSTM, GradCAM as VideoGradCAM, overlay_cam  # noqa: E402

try:
    from neurosymbolic import run_neurosymbolic_assessment  # type: ignore  # noqa: E402
except Exception:
    run_neurosymbolic_assessment = None

PRIMARY_WEIGHTS_PATH = ROOT / 'models_adv' / 'best_model_weights.pt'
FALLBACK_WEIGHTS_PATH = ROOT / 'integration_assets' / 'best_model_weights.pt'
WEIGHTS_PATH = PRIMARY_WEIGHTS_PATH if PRIMARY_WEIGHTS_PATH.exists() else FALLBACK_WEIGHTS_PATH
GAN_DIFF_WEIGHTS_PATH = ROOT / 'models_gan_vs_diffusion' / 'best_model_weights.pt'
GAN_DIFF_CONFIG_PATH = ROOT / 'models_gan_vs_diffusion' / 'config.json'
VIDEO_WEIGHTS_PATH = VIDEO_CODE_DIR / 'video_xception_lstm.pt'
DEVICE = torch.device('cpu')
IMAGE_SIZE = 224
THRESHOLD_AI = 0.50
THRESHOLD_SUSPECT = 0.35
VIDEO_THRESHOLD_AI = 0.48
VIDEO_THRESHOLD_SUSPECT = 0.30


class FFTFeatureExtractor(nn.Module):
    def __init__(self, output_dim: int = 512):
        super().__init__()
        self.fft_processor = nn.Sequential(
            nn.Linear(12, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1),
            nn.Linear(64, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Linear(128, output_dim),
        )

    @torch.no_grad()
    def _extract_fft_features(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, channels, height, width = x.shape
        x_f32 = x.float()
        if channels == 3:
            gray = 0.299 * x_f32[:, 0] + 0.587 * x_f32[:, 1] + 0.114 * x_f32[:, 2]
        else:
            gray = x_f32[:, 0]
        fft_img = torch.fft.fft2(gray)
        fft_shift = torch.fft.fftshift(fft_img)
        magnitude = torch.abs(fft_shift) + 1e-8
        magnitude = magnitude / (
            magnitude.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0] + 1e-8
        )

        fft_features = []
        for index in range(batch_size):
            mag = magnitude[index].flatten()
            feature_vector = torch.stack(
                [
                    mag.mean(),
                    mag.std().clamp(min=1e-8),
                    mag.max(),
                    mag.min(),
                    (mag > mag.mean()).float().mean(),
                    mag.median(),
                    magnitude[index][: height // 4, :].mean(),
                    magnitude[index][height // 4 : height // 2, :].mean(),
                    magnitude[index][height // 2 : 3 * height // 4, :].mean(),
                    magnitude[index][3 * height // 4 :, :].mean(),
                    (mag > 0.5).float().mean(),
                    (mag > 0.1).float().mean(),
                ]
            )
            fft_features.append(torch.clamp(feature_vector, min=-10, max=10))

        return torch.stack(fft_features, dim=0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        fft_features = self._extract_fft_features(x)
        fft_features = fft_features.to(x.dtype).detach()
        return self.fft_processor(fft_features)


class EfficientNetFFTFusion(nn.Module):
    def __init__(self, num_classes: int = 2, dropout: float = 0.4, backbone: str = 'efficientnet_b2'):
        super().__init__()
        self.backbone = timm.create_model(backbone, pretrained=False, num_classes=0)
        backbone_dim = self.backbone.num_features
        fft_dim = 512
        self.fft_extractor = FFTFeatureExtractor(output_dim=fft_dim)
        fusion_dim = backbone_dim + fft_dim
        self.fusion = nn.Sequential(
            nn.Linear(fusion_dim, 1024),
            nn.LayerNorm(1024),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(1024, 512),
            nn.LayerNorm(512),
            nn.GELU(),
            nn.Dropout(dropout * 0.5),
        )
        self.classifier = nn.Linear(512, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        backbone_features = self.backbone(x)
        fft_features = self.fft_extractor(x)
        fused = torch.cat([backbone_features, fft_features], dim=1)
        fused = self.fusion(fused)
        return self.classifier(fused)


@dataclass
class InferenceArtifacts:
    probability_ai: float
    probability_authentic: float
    verdict: str
    confidence_score: float
    heatmap_regions: list[dict[str, Any]]
    fft_summary: dict[str, Any]
    metadata: dict[str, Any]
    model_breakdown: list[dict[str, Any]]
    artifacts: list[dict[str, Any]]
    source_analysis: dict[str, Any] | None
    ethical: dict[str, Any] | None
    neurosymbolic: dict[str, Any] | None


@dataclass
class SimpleGradCAM:
    model: nn.Module
    target_layer: nn.Module
    activations: list[torch.Tensor]
    gradients: list[torch.Tensor]

    def __init__(self, model: nn.Module, target_layer: nn.Module):
        self.model = model
        self.target_layer = target_layer
        self.activations = []
        self.gradients = []
        self.forward_handle = target_layer.register_forward_hook(self._save_activation)
        self.backward_handle = target_layer.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module, inputs, output):
        self.activations = [output]

    def _save_gradient(self, module, grad_input, grad_output):
        self.gradients = [grad_output[0]]

    def generate_cam(self, input_tensor: torch.Tensor, target_class: int = 1) -> np.ndarray:
        self.model.eval()
        output = self.model(input_tensor)
        self.model.zero_grad()
        output[0, target_class].backward(retain_graph=True)
        activation = self.activations[0].detach()
        gradient = self.gradients[0].detach()
        weights = gradient.mean(dim=(2, 3), keepdim=True)
        cam = (weights * activation).sum(dim=1, keepdim=True)
        cam = F.relu(cam)
        cam = cam - cam.min()
        cam = cam / (cam.max() + 1e-8)
        cam = F.interpolate(cam, size=(IMAGE_SIZE, IMAGE_SIZE), mode='bilinear', align_corners=False)
        return cam.squeeze().cpu().numpy()

    def close(self):
        self.forward_handle.remove()
        self.backward_handle.remove()


def detect_backbone(state_dict: dict[str, torch.Tensor]) -> str:
    fusion_in_dim = state_dict['fusion.0.weight'].shape[1]
    backbone_dim = fusion_in_dim - 512
    backbone_map = {1280: 'efficientnet_b0', 1408: 'efficientnet_b2', 1792: 'efficientnet_b4'}
    return backbone_map.get(backbone_dim, 'efficientnet_b2')


state_dict = torch.load(WEIGHTS_PATH, map_location=DEVICE)
model = EfficientNetFFTFusion(backbone=detect_backbone(state_dict))
model.load_state_dict(state_dict, strict=False)
for parameter in model.parameters():
    parameter.requires_grad_(True)
model = model.to(DEVICE)
model.eval()
MODEL_INFO = {'model_type': 'efficientnet_fft', 'backbone': detect_backbone(state_dict), 'optimal_threshold': THRESHOLD_AI}

GAN_DIFF_MODEL = None
GAN_DIFF_CONFIG = None
if GAN_DIFF_WEIGHTS_PATH.exists() and GAN_DIFF_CONFIG_PATH.exists():
    GAN_DIFF_CONFIG = json.loads(GAN_DIFF_CONFIG_PATH.read_text())
    gan_diff_model = models.resnet18(weights=None)
    gan_diff_model.fc = nn.Linear(gan_diff_model.fc.in_features, 2)
    gan_diff_state = torch.load(GAN_DIFF_WEIGHTS_PATH, map_location=DEVICE)
    if isinstance(gan_diff_state, dict) and 'model_state' in gan_diff_state:
        gan_diff_state = gan_diff_state['model_state']
    gan_diff_model.load_state_dict(gan_diff_state, strict=False)
    GAN_DIFF_MODEL = gan_diff_model.to(DEVICE)
    GAN_DIFF_MODEL.eval()

VIDEO_MODEL = None
VIDEO_CONFIG = None
if VIDEO_WEIGHTS_PATH.exists():
    video_ckpt = torch.load(VIDEO_WEIGHTS_PATH, map_location=DEVICE)
    VIDEO_CONFIG = video_ckpt.get('config', {})
    VIDEO_MODEL = ResNetLSTM(
        hidden_size=VIDEO_CONFIG.get('hidden_size', 256),
        num_layers=VIDEO_CONFIG.get('num_layers', 1),
        bidirectional=VIDEO_CONFIG.get('bidirectional', True),
        temporal_pool=VIDEO_CONFIG.get('temporal_pool', 'mean'),
        pretrained=False,
        backbone_name=VIDEO_CONFIG.get('backbone', 'xception'),
    )
    VIDEO_MODEL.load_state_dict(video_ckpt['model_state'], strict=False)
    for parameter in VIDEO_MODEL.parameters():
        parameter.requires_grad_(True)
    VIDEO_MODEL = VIDEO_MODEL.to(DEVICE)
    VIDEO_MODEL.eval()


def pad_to_min_size(image: Image.Image, size: int) -> Image.Image:
    width, height = image.size
    pad_w = max(0, size - width)
    pad_h = max(0, size - height)
    if not (pad_w or pad_h):
        return image
    left = pad_w // 2
    top = pad_h // 2
    right = pad_w - left
    bottom = pad_h - top
    arr = np.array(image)
    arr = np.pad(arr, ((top, bottom), (left, right), (0, 0)), mode='reflect')
    return Image.fromarray(arr.astype(np.uint8))


transform = transforms.Compose([
    transforms.Lambda(lambda image: pad_to_min_size(image, IMAGE_SIZE)),
    transforms.CenterCrop(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

video_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

app = FastAPI(title='UAIDE Inference API', version='1.3.0')
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)

FRONTEND_DIST_DIR = ROOT / 'dist'
FRONTEND_ASSETS_DIR = FRONTEND_DIST_DIR / 'assets'


def image_to_base64(image: Image.Image) -> str:
    buf = io.BytesIO()
    image.save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode('utf-8')


def softmax(logits: torch.Tensor) -> torch.Tensor:
    return torch.softmax(logits, dim=1)


def classify_verdict(probability_ai: float, *, is_video: bool = False) -> tuple[str, float]:
    ai_th = VIDEO_THRESHOLD_AI if is_video else THRESHOLD_AI
    suspect_th = VIDEO_THRESHOLD_SUSPECT if is_video else THRESHOLD_SUSPECT
    if probability_ai >= ai_th:
        return 'ai_generated', probability_ai * 100
    if probability_ai >= suspect_th:
        margin = (probability_ai - suspect_th) / max(ai_th - suspect_th, 1e-6)
        return 'suspect', (55 + margin * 25)
    return 'authentic', (1 - probability_ai) * 100


def analyze_fft(image: Image.Image) -> dict[str, Any]:
    gray = np.asarray(image.convert('L'), dtype=np.float32) / 255.0
    spectrum = np.fft.fftshift(np.fft.fft2(gray))
    magnitude = np.log1p(np.abs(spectrum))
    h, w = magnitude.shape
    y, x = np.indices((h, w))
    center_y, center_x = h / 2.0, w / 2.0
    radius = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2)
    max_radius = float(radius.max() or 1.0)

    def band_energy(start_ratio: float, end_ratio: float) -> float:
        mask = (radius >= max_radius * start_ratio) & (radius < max_radius * end_ratio)
        return float(magnitude[mask].mean()) if np.any(mask) else 0.0

    low = band_energy(0.0, 0.2)
    mid = band_energy(0.2, 0.55)
    high = band_energy(0.55, 1.0)
    peak = float(magnitude.max())
    anomaly = high > low * 1.18
    anomaly_bands = []
    if mid > low * 1.05:
        anomaly_bands.append('mid-band elevation')
    if high > low * 1.18:
        anomaly_bands.append('high-frequency grid energy')
    if not anomaly_bands:
        anomaly_bands.append('no major spectral spikes')

    return {
        'peakFrequency': f'{peak:.2f} spectral units',
        'spectralAnomaly': anomaly,
        'anomalyBands': anomaly_bands,
        'dctCoefficients': f'Low {low:.3f} · Mid {mid:.3f} · High {high:.3f}',
        'noisePattern': 'Periodic upsampling artifact suspected' if anomaly else 'Natural broadband texture distribution',
        'bands': {'low': low, 'mid': mid, 'high': high},
    }


def overlay_gradcam_from_pil(pil_image: Image.Image) -> tuple[Image.Image, np.ndarray]:
    model.eval()
    for parameter in model.parameters():
        parameter.requires_grad_(True)
    input_tensor = transform(pil_image).unsqueeze(0).to(DEVICE)
    input_tensor.requires_grad_(True)
    target_layer = model.backbone.conv_head
    grad_cam = SimpleGradCAM(model, target_layer)
    try:
        model.zero_grad(set_to_none=True)
        cam = grad_cam.generate_cam(input_tensor, target_class=1)
    finally:
        grad_cam.close()

    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
    original = cv2.resize(np.array(pil_image.convert('RGB')), (IMAGE_SIZE, IMAGE_SIZE))
    overlay = cv2.addWeighted(original, 0.6, heatmap, 0.4, 0)
    return Image.fromarray(overlay), cam


def cam_to_regions(cam: np.ndarray) -> list[dict[str, Any]]:
    regions = []
    rows = cols = 4
    h, w = cam.shape
    tile_h = h / rows
    tile_w = w / cols
    scored = []
    for row in range(rows):
        for col in range(cols):
            y1 = int(row * tile_h)
            y2 = int((row + 1) * tile_h)
            x1 = int(col * tile_w)
            x2 = int((col + 1) * tile_w)
            score = float(cam[y1:y2, x1:x2].mean())
            scored.append((score, row, col))
    scored.sort(reverse=True, key=lambda item: item[0])
    labels = ['Primary anomaly', 'Boundary inconsistency', 'Texture smoothing', 'Frequency spike']
    for index, (score, row, col) in enumerate(scored[:4]):
        regions.append({
            'x': round((col * 100) / cols, 2),
            'y': round((row * 100) / rows, 2),
            'w': round(100 / cols, 2),
            'h': round(100 / rows, 2),
            'intensity': round(float(score), 3),
            'label': labels[index] if index < len(labels) else f'Region {index + 1}',
        })
    return regions


def collect_metadata(image: Image.Image, filename: str, content_type: str, size_bytes: int) -> dict[str, Any]:
    stat = ImageStat.Stat(image.convert('RGB'))
    exif = image.getexif() if hasattr(image, 'getexif') else None
    return {
        'mimeType': content_type,
        'dimensions': f'{image.width} × {image.height} px',
        'channels': 'RGB',
        'meanRgb': ', '.join(f'{value:.1f}' for value in stat.mean),
        'stdRgb': ', '.join(f'{value:.1f}' for value in stat.stddev),
        'bitDepth': '8-bit',
        'fileName': filename,
        'fileSizeBytes': str(size_bytes),
        'softwareTag': str(exif.get(305, 'Not present')) if exif else 'Not present',
        'cameraModel': str(exif.get(272, 'Not present')) if exif else 'Not present',
        'creationDate': str(exif.get(306, 'Not present')) if exif else 'Not present',
        'exifFields': str(len(exif)) if exif else '0',
    }


def predict_source_from_pil(image: Image.Image) -> dict[str, Any] | None:
    if GAN_DIFF_MODEL is None or GAN_DIFF_CONFIG is None:
        return None
    image_size = int(GAN_DIFF_CONFIG.get('image_size', 224))
    id_to_label = GAN_DIFF_CONFIG.get('id_to_label', {0: 'gan', 1: 'diffusion'})
    source_transform = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    with torch.no_grad():
        input_tensor = source_transform(image).unsqueeze(0).to(DEVICE)
        logits = GAN_DIFF_MODEL(input_tensor)
        probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
        pred_idx = int(np.argmax(probs))
    labels = {int(k): str(v) for k, v in id_to_label.items()} if isinstance(id_to_label, dict) else {0: 'gan', 1: 'diffusion'}
    predicted_source = labels.get(pred_idx, 'unknown')
    top_prob = float(probs[pred_idx])
    return {
        'predictedSource': predicted_source,
        'label': predicted_source,
        'top_prob': round(top_prob, 6),
        'ganProbability': round(float(probs[0]), 6),
        'diffusionProbability': round(float(probs[1]), 6),
    }


def build_artifacts(probability_ai: float, fft_summary: dict[str, Any], metadata: dict[str, Any], regions: list[dict[str, Any]], ethical: dict[str, Any] | None) -> list[dict[str, Any]]:
    highest = max((region['intensity'] for region in regions), default=0.0)
    artifacts = [
        {'id': 1, 'type': 'Fusion Model Confidence', 'severity': 'critical' if probability_ai >= THRESHOLD_AI else 'medium' if probability_ai >= THRESHOLD_SUSPECT else 'low', 'detail': f'EfficientNet-B2 + FFT fusion scored AI likelihood at {probability_ai * 100:.2f}%.'},
        {'id': 2, 'type': 'Frequency Domain Signal', 'severity': 'high' if fft_summary['spectralAnomaly'] else 'low', 'detail': fft_summary['noisePattern']},
        {'id': 3, 'type': 'Localized Artifact Region', 'severity': 'high' if highest > 0.72 else 'medium', 'detail': f'Top suspicious tile intensity measured at {highest * 100:.1f}%.'},
        {'id': 4, 'type': 'Metadata Audit', 'severity': 'medium' if metadata['cameraModel'] == 'Not present' else 'low', 'detail': f"Camera model: {metadata['cameraModel']}; software tag: {metadata['softwareTag']}."},
    ]
    if ethical:
        artifacts.append({'id': 5, 'type': 'Ethical Assessment', 'severity': 'critical' if not ethical['is_ethical'] else 'low', 'detail': ethical['simpleStatus']})
    return artifacts


def run_ethical_assessment(image_array: np.ndarray) -> dict[str, Any]:
    assessment = EthicalAssessment.assess(image_array, threshold=0.5)
    return {
        'is_ethical': bool(assessment.get('is_ethical', False)),
        'status': assessment.get('status', 'UNKNOWN'),
        'riskScore': float(assessment.get('risk_score', 0.0)),
        'flags': assessment.get('flags', []),
        'simpleStatus': get_simple_status(assessment),
        'report': format_ethical_report(assessment),
    }


def extract_video_frames(video_path: str, frames_per_video: int = 16, frame_stride: int = 4) -> tuple[list[np.ndarray], float, int]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError('Unable to open video file')
    fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    duration = total_frames / fps if fps else 0

    target_frames = max(int(frames_per_video), 16)
    stride = max(int(frame_stride), 1)
    frames = []

    if total_frames > 0:
        candidate_indices = np.arange(0, total_frames, stride, dtype=int)
        if len(candidate_indices) > target_frames:
            sample_positions = np.linspace(0, len(candidate_indices) - 1, num=target_frames, dtype=int)
            candidate_indices = candidate_indices[sample_positions]

        for frame_index in candidate_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
            ok, frame = cap.read()
            if not ok:
                continue
            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    else:
        idx = 0
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            if idx % stride == 0:
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                if len(frames) >= target_frames:
                    break
            idx += 1

    cap.release()
    if not frames:
        raise RuntimeError('No frames extracted from video')
    return frames, duration, total_frames


def analyze_video_file(video_path: str, filename: str, content_type: str, size_bytes: int) -> dict[str, Any]:
    if VIDEO_MODEL is None or VIDEO_CONFIG is None:
        raise HTTPException(status_code=501, detail='Video model not integrated yet.')

    frames, duration, total_frames = extract_video_frames(
        video_path,
        frames_per_video=VIDEO_CONFIG.get('frames_per_video', 16),
        frame_stride=VIDEO_CONFIG.get('frame_stride', 4),
    )
    pil_frames = [Image.fromarray(frame) for frame in frames]
    frame_tensors = [video_transform(frame).unsqueeze(0) for frame in pil_frames]
    clip = torch.stack([tensor.squeeze(0) for tensor in frame_tensors], dim=0).unsqueeze(0).to(DEVICE)
    clip.requires_grad_(True)

    with torch.no_grad():
        frame_logits, video_logits = VIDEO_MODEL(clip)
        frame_probs = torch.softmax(frame_logits.squeeze(0), dim=1)[:, 1].cpu().numpy()
        video_probs = torch.softmax(video_logits, dim=1)[0].cpu().numpy()
    prob_fake = float(video_probs[1])
    frame_mean = float(frame_probs.mean()) if len(frame_probs) else prob_fake
    frame_peak = float(frame_probs.max()) if len(frame_probs) else prob_fake
    temporal_consistency = min(1.0, max(0.0, frame_peak * 0.55 + frame_mean * 0.45))
    fused_video_score = float(min(1.0, prob_fake * 0.55 + frame_mean * 0.25 + frame_peak * 0.20))
    verdict, confidence_score = classify_verdict(fused_video_score, is_video=True)

    lead_idx = int(np.argmax(frame_probs))
    lead_frame = frames[lead_idx]
    target_layer = VIDEO_MODEL.get_gradcam_target_layer()
    grad_cam = VideoGradCAM(VIDEO_MODEL, target_layer)
    cam = grad_cam.generate(clip[:, lead_idx:lead_idx + 1], class_idx=1)
    overlay = overlay_cam(lead_frame, cam, alpha=0.45)
    overlay_image = Image.fromarray(overlay)
    heatmap_regions = cam_to_regions(cv2.resize(cam, (224, 224)))

    flagged_segments = []
    clean_segments = []
    fps_est = total_frames / duration if duration > 0 else 24.0
    segment_length = duration / max(len(frame_probs), 1) if duration > 0 else 1.0
    current_clean_start = 0.0
    for index, score in enumerate(frame_probs):
        start = index * segment_length
        end = min(duration, (index + 1) * segment_length)
        effective_score = max(float(score), fused_video_score * 0.75)
        if effective_score >= VIDEO_THRESHOLD_SUSPECT:
            severity = 'critical' if effective_score >= VIDEO_THRESHOLD_AI else 'medium'
            flagged_segments.append({
                'start': round(start, 2),
                'end': round(end, 2),
                'severity': severity,
                'reason': 'Temporal manipulation spike',
                'frames': [int(index * VIDEO_CONFIG.get('frame_stride', 4))],
            })
        else:
            clean_segments.append({'start': round(start, 2), 'end': round(end, 2)})

    lead_image = pil_frames[lead_idx]
    ethical = run_ethical_assessment(np.array(lead_image.convert('RGB'))) if verdict != 'authentic' else None
    source_analysis = predict_source_from_pil(lead_image) if verdict != 'authentic' else None
    neurosymbolic = None
    if verdict != 'authentic' and run_neurosymbolic_assessment is not None:
        try:
            neurosymbolic = run_neurosymbolic_assessment(
                np.array(lead_image.convert('RGB')),
                fused_video_score,
                source_analysis,
                {
                    'risk_score': ethical['riskScore'],
                    'status': ethical['status'],
                } if ethical else None,
            )
        except Exception:
            neurosymbolic = None
    metadata = {
        'mimeType': content_type,
        'dimensions': f'{lead_image.width} × {lead_image.height} px',
        'durationSeconds': round(duration, 2),
        'frameCount': total_frames,
        'sampledFrames': len(frames),
        'frameStride': VIDEO_CONFIG.get('frame_stride', 4),
        'backbone': VIDEO_CONFIG.get('backbone', 'xception'),
        'temporalMeanScore': round(frame_mean, 6),
        'temporalPeakScore': round(frame_peak, 6),
        'fusedVideoScore': round(fused_video_score, 6),
        'temporalConsistency': round(temporal_consistency, 6),
        'heatmapPreview': image_to_base64(overlay_image),
    }

    model_breakdown = [
        {'model': f"Video {VIDEO_CONFIG.get('backbone', 'xception').upper()} + LSTM", 'score': round(prob_fake * 100, 2), 'weight': 0.55},
        {'model': 'Frame anomaly mean', 'score': round(frame_mean * 100, 2), 'weight': 0.25},
        {'model': 'Peak frame anomaly', 'score': round(frame_peak * 100, 2), 'weight': 0.20},
    ]
    if source_analysis:
        model_breakdown.append({'model': f"AI Source: {source_analysis['predictedSource'].upper()}", 'score': round(max(source_analysis['ganProbability'], source_analysis['diffusionProbability']) * 100, 2), 'weight': 0.2})

    artifacts = [
        {'id': 1, 'type': 'Temporal Inconsistency', 'severity': 'critical' if fused_video_score >= VIDEO_THRESHOLD_AI else 'medium', 'detail': f'Fused video AI likelihood scored {fused_video_score * 100:.2f}% (sequence model raw: {prob_fake * 100:.2f}%).'},
        {'id': 2, 'type': 'Flagged Frames', 'severity': 'medium', 'detail': f'{len(flagged_segments)} suspicious segments identified across {len(frames)} sampled frames.'},
    ]
    if ethical:
        artifacts.append({'id': 3, 'type': 'Ethical Assessment', 'severity': 'critical' if not ethical['is_ethical'] else 'low', 'detail': ethical['simpleStatus']})

    return {
        'type': 'video',
        'filename': filename,
        'filesize': f'{size_bytes / (1024 * 1024):.2f} MB',
        'resolution': f'{lead_image.width} × {lead_image.height} px',
        'format': 'MP4',
        'duration': time.strftime('%H:%M:%S', time.gmtime(duration)),
        'frameRate': f'{fps_est:.0f} fps',
        'totalFrames': total_frames,
        'analysisId': f"UAD-{uuid.uuid4().hex[:10].upper()}",
        'processingTime': 'Video inference complete',
        'verdict': verdict,
        'confidenceScore': round(confidence_score, 2),
        'modelBreakdown': model_breakdown,
        'gradcam': {'regions': heatmap_regions},
        'timeline': {'flaggedSegments': flagged_segments, 'cleanSegments': clean_segments},
        'artifacts': artifacts,
        'fft': {'peakFrequency': 'Temporal model', 'spectralAnomaly': len(flagged_segments) > 0, 'anomalyBands': ['temporal coherence'], 'dctCoefficients': 'Frame-level sequence modelling', 'noisePattern': 'Temporal inconsistency tracking', 'bands': {'low': 0.0, 'mid': 0.0, 'high': 0.0}},
        'metadata': metadata,
        'sourceAnalysis': source_analysis,
        'ethical': ethical,
        'neurosymbolic': neurosymbolic,
        'raw': {
            'probabilityAi': round(fused_video_score, 6),
            'rawSequenceProbabilityAi': round(prob_fake, 6),
            'frameMeanProbabilityAi': round(frame_mean, 6),
            'framePeakProbabilityAi': round(frame_peak, 6),
            'probabilityAuthentic': round(float(video_probs[0]), 6),
            'thresholds': {'suspect': VIDEO_THRESHOLD_SUSPECT, 'aiGenerated': VIDEO_THRESHOLD_AI},
            'ganDiffusionWeightsLoaded': GAN_DIFF_MODEL is not None,
            'videoWeightsLoaded': VIDEO_MODEL is not None,
        },
    }


def format_result(image: Image.Image, filename: str, content_type: str, size_bytes: int) -> InferenceArtifacts:
    tensor = transform(image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        logits = model(tensor)
        probabilities = softmax(logits)[0].cpu().numpy()
    probability_authentic = float(probabilities[0])
    probability_ai = float(probabilities[1])
    verdict, confidence_score = classify_verdict(probability_ai)
    fft_summary = analyze_fft(image)
    overlay_image, cam = overlay_gradcam_from_pil(image)
    heatmap_regions = cam_to_regions(cam)
    metadata = collect_metadata(image, filename, content_type, size_bytes)
    ethical = run_ethical_assessment(np.array(image.convert('RGB'))) if verdict != 'authentic' else None
    artifacts = build_artifacts(probability_ai, fft_summary, metadata, heatmap_regions, ethical)
    source_analysis = predict_source_from_pil(image) if verdict != 'authentic' else None
    neurosymbolic = None
    if verdict != 'authentic' and run_neurosymbolic_assessment is not None:
        try:
            neurosymbolic = run_neurosymbolic_assessment(
                np.array(image.convert('RGB')),
                probability_ai,
                source_analysis,
                {
                    'risk_score': ethical['riskScore'],
                    'status': ethical['status'],
                } if ethical else None,
            )
        except Exception:
            neurosymbolic = None
    fft_branch_score = round(min(99.9, fft_summary['bands']['high'] / max(fft_summary['bands']['low'], 1e-6) * 40), 2)
    model_breakdown = [
        {'model': 'EfficientNet-B2 Spatial', 'score': round(probability_ai * 100, 2), 'weight': 0.55},
        {'model': 'FFT Frequency Branch', 'score': fft_branch_score, 'weight': 0.25},
    ]
    if source_analysis:
        model_breakdown.append({'model': f"AI Source: {source_analysis['predictedSource'].upper()}", 'score': round(max(source_analysis['ganProbability'], source_analysis['diffusionProbability']) * 100, 2), 'weight': 0.20})
    metadata['heatmapPreview'] = image_to_base64(overlay_image)
    return InferenceArtifacts(probability_ai, probability_authentic, verdict, round(confidence_score, 2), heatmap_regions, fft_summary, metadata, model_breakdown, artifacts, source_analysis, ethical, neurosymbolic)


@app.get('/api/health')
def healthcheck() -> dict[str, Any]:
    return {'ok': True, 'modelLoaded': WEIGHTS_PATH.exists(), 'model': 'EfficientNet-B2 + FFT Fusion', 'ganDiffusionLoaded': GAN_DIFF_MODEL is not None, 'videoModelLoaded': VIDEO_MODEL is not None}



if FRONTEND_ASSETS_DIR.exists():
    app.mount('/assets', StaticFiles(directory=str(FRONTEND_ASSETS_DIR)), name='assets')


@app.get('/', include_in_schema=False)
def serve_index() -> FileResponse:
    index_file = FRONTEND_DIST_DIR / 'index.html'
    if not index_file.exists():
        raise HTTPException(status_code=503, detail='Frontend build not found')
    return FileResponse(index_file)


@app.get('/{full_path:path}', include_in_schema=False)
def serve_spa(full_path: str):
    if full_path.startswith('api/'):
        raise HTTPException(status_code=404, detail='Not found')
    candidate = FRONTEND_DIST_DIR / full_path
    if candidate.exists() and candidate.is_file():
        return FileResponse(candidate)
    index_file = FRONTEND_DIST_DIR / 'index.html'
    if not index_file.exists():
        raise HTTPException(status_code=503, detail='Frontend build not found')
    return FileResponse(index_file)


@app.post('/api/analyze')
async def analyze_media(file: UploadFile = File(...)) -> JSONResponse:
    inferred_content_type, _ = mimetypes.guess_type(file.filename or '')
    content_type = file.content_type or inferred_content_type or ''
    if content_type == 'application/octet-stream' and inferred_content_type:
        content_type = inferred_content_type

    if not content_type.startswith(('image/', 'video/')):
        raise HTTPException(status_code=400, detail='Unsupported media type')
    contents = await file.read()
    if not contents:
        raise HTTPException(status_code=400, detail='Empty file received')

    if content_type.startswith('video/'):
        suffix = Path(file.filename or 'upload.mp4').suffix or '.mp4'
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir='/home/user/app/backend') as tmp:
            tmp.write(contents)
            video_path = tmp.name
        try:
            payload = analyze_video_file(video_path, file.filename or 'upload', content_type, len(contents))
            return JSONResponse(payload)
        finally:
            try:
                os.remove(video_path)
            except OSError:
                pass

    start = time.perf_counter()
    try:
        image = Image.open(io.BytesIO(contents)).convert('RGB')
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f'Unable to open image: {exc}') from exc

    artifacts = format_result(image, file.filename or 'upload', content_type, len(contents))
    elapsed = time.perf_counter() - start
    payload = {
        'type': 'image',
        'filename': file.filename or 'upload',
        'filesize': f'{len(contents) / (1024 * 1024):.2f} MB',
        'resolution': f'{image.width} × {image.height} px',
        'format': image.format or (content_type.split('/')[-1].upper()),
        'analysisId': f"UAD-{uuid.uuid4().hex[:10].upper()}",
        'processingTime': f'{elapsed:.2f}s',
        'verdict': artifacts.verdict,
        'confidenceScore': artifacts.confidence_score,
        'modelBreakdown': artifacts.model_breakdown,
        'gradcam': {'regions': artifacts.heatmap_regions},
        'artifacts': artifacts.artifacts,
        'fft': artifacts.fft_summary,
        'metadata': artifacts.metadata,
        'sourceAnalysis': artifacts.source_analysis,
        'ethical': artifacts.ethical,
        'neurosymbolic': artifacts.neurosymbolic,
        'raw': {
            'probabilityAi': round(artifacts.probability_ai, 6),
            'probabilityAuthentic': round(artifacts.probability_authentic, 6),
            'thresholds': {'suspect': THRESHOLD_SUSPECT, 'aiGenerated': THRESHOLD_AI},
            'weightsPath': str(WEIGHTS_PATH.name),
            'ganDiffusionWeightsLoaded': GAN_DIFF_MODEL is not None,
            'videoWeightsLoaded': VIDEO_MODEL is not None,
        },
    }
    return JSONResponse(payload)