main.py

import os
import io
import math
import uuid
import shutil
import asyncio
import threading
from datetime import datetime, timedelta
from functools import partial
from pathlib import Path
from typing import List, Dict, Any, Optional

import cv2
import numpy as np
import torch
import uvicorn
import httpx
from huggingface_hub import hf_hub_download

from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from MyModel import PollutionDifferenceModel
from my_Segmenter import Segmenter

import json
from collections import defaultdict


# =========================
# 全局进度存储
# =========================
batch_progress = defaultdict(dict)
progress_lock = threading.Lock()


# =========================
# 基础目录初始化
# =========================
BASE_DIR = Path(".")
STATIC_DIR = BASE_DIR / "static"
STATIC_RESULTS_DIR = STATIC_DIR / "results"
RUNS_DIR = BASE_DIR / "runs"
AIR_STATION_DIR = BASE_DIR / "AirStationImage"
AIR_POLLUTION_MAPPING_DIR = BASE_DIR / "air pollution mapping"
FRONTEND_DIR = BASE_DIR / "frontend"

MAPILLARY_ACCESS_TOKEN = os.getenv("MAPILLARY_ACCESS_TOKEN", "").strip()
MAPILLARY_API = "https://graph.mapillary.com"

for d in [STATIC_DIR, STATIC_RESULTS_DIR, RUNS_DIR, AIR_STATION_DIR, AIR_POLLUTION_MAPPING_DIR]:
    d.mkdir(parents=True, exist_ok=True)


# =========================
# Hugging Face private model repo 配置
# =========================
HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "").strip()
HF_TOKEN = os.getenv("HF_TOKEN", "").strip() or None

MODEL_FILENAMES = {
    "CO":   "best_CO_model_multiscale20251110.pth",
    "NO2":  "best_NO2_model_multiscale20251110.pth",
    "PM25": "best_PM25_model_multiscale20251110.pth",
    "PM10": "best_PM10_model_multiscale20251110.pth",
    "O3":   "best_O3_model_multiscale20251110.pth",
}

loaded_models: Dict[str, PollutionDifferenceModel] = {}
model_lock = threading.Lock()


# =========================
# 文件上传限制
# =========================
MAX_FILE_SIZE = 10 * 1024 * 1024
ALLOWED_CONTENT_TYPES = {"image/jpeg", "image/png", "image/webp"}


# =========================
# 污染物合理范围校验
# =========================
POLLUTANT_RANGES = {
    "CO":   (0, 50),
    "NO2":  (0, 500),
    "PM25": (0, 999),
    "PM10": (0, 999),
    "O3":   (0, 500),
}


# =========================
# 初始化分割模型
# =========================
segmenter = Segmenter(dataset="cityscapes", task="semantic", device="cpu")


# =========================
# FastAPI 初始化
# =========================
app = FastAPI(title="香港空气污染预测")

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

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
app.mount("/runs", StaticFiles(directory=str(RUNS_DIR)), name="runs")
app.mount("/AirStationImage", StaticFiles(directory=str(AIR_STATION_DIR)), name="AirStationImage")
app.mount(
    "/air-pollution-mapping",
    StaticFiles(directory=str(AIR_POLLUTION_MAPPING_DIR)),
    name="air-pollution-mapping"
)


def haversine_m(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    radius = 6_371_000.0
    p1, p2 = math.radians(lat1), math.radians(lat2)
    dp = math.radians(lat2 - lat1)
    dl = math.radians(lon2 - lon1)
    value = math.sin(dp / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
    return 2 * radius * math.asin(math.sqrt(value))


def panorama_view(
    image: np.ndarray,
    center_x: float,
    center_y: float,
    fov_deg: float = 90.0,
    width: int = 768,
    height: int = 512
) -> np.ndarray:
    focal = width / (2 * math.tan(math.radians(fov_deg) / 2))
    px, py = np.meshgrid(np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32))

    x = (px - (width - 1) / 2) / focal
    y = -((py - (height - 1) / 2) / focal)
    z = np.ones_like(x)

    norm = np.sqrt(x * x + y * y + z * z)
    x, y, z = x / norm, y / norm, z / norm

    yaw = math.radians((center_x - 0.5) * 360.0)
    pitch = math.radians(max(-75.0, min(75.0, (0.5 - center_y) * 180.0)))

    cp, sp = math.cos(pitch), math.sin(pitch)
    y2, z2 = y * cp + z * sp, -y * sp + z * cp

    cy, sy = math.cos(yaw), math.sin(yaw)
    x3, z3 = x * cy + z2 * sy, -x * sy + z2 * cy

    longitude = np.arctan2(x3, z3)
    latitude = np.arcsin(np.clip(y2, -1, 1))

    source_h, source_w = image.shape[:2]
    map_x = ((longitude / (2 * np.pi) + 0.5) * source_w).astype(np.float32)
    map_y = np.clip((0.5 - latitude / np.pi) * source_h, 0, source_h - 1).astype(np.float32)

    return cv2.remap(image, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP)


async def mapillary_get(path: str, params: Dict[str, Any]) -> Dict[str, Any]:
    if not MAPILLARY_ACCESS_TOKEN:
        raise HTTPException(status_code=503, detail="Mapillary 尚未設定。請配置 MAPILLARY_ACCESS_TOKEN。")

    request_params = {**params, "access_token": MAPILLARY_ACCESS_TOKEN}

    try:
        async with httpx.AsyncClient(timeout=20, follow_redirects=True) as client:
            response = await client.get(f"{MAPILLARY_API}/{path.lstrip('/')}", params=request_params)

        response.raise_for_status()
        return response.json()

    except httpx.HTTPStatusError as exc:
        raise HTTPException(
            status_code=502,
            detail=f"Mapillary API 回應錯誤（HTTP {exc.response.status_code}）。"
        ) from exc

    except (httpx.HTTPError, ValueError) as exc:
        raise HTTPException(status_code=502, detail="暫時無法連接 Mapillary，請稍後再試。") from exc


def get_model_path_from_hf(pollutant: str) -> str:
    if pollutant not in MODEL_FILENAMES:
        raise HTTPException(status_code=400, detail=f"不支持的污染物类型: {pollutant}")

    if not HF_MODEL_REPO:
        raise HTTPException(
            status_code=500,
            detail="未设置 HF_MODEL_REPO。请在 Hugging Face Space 的 Secrets 中配置 HF_MODEL_REPO。"
        )

    filename = MODEL_FILENAMES[pollutant]

    try:
        return hf_hub_download(
            repo_id=HF_MODEL_REPO,
            filename=filename,
            repo_type="model",
            token=HF_TOKEN,
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"无法从 Hugging Face 下载模型文件 {filename}: {str(e)}"
        )


def load_pollution_model(pollutant: str) -> PollutionDifferenceModel:
    if pollutant not in MODEL_FILENAMES:
        raise HTTPException(status_code=400, detail=f"不支持的污染物类型: {pollutant}")

    if pollutant in loaded_models:
        return loaded_models[pollutant]

    with model_lock:
        if pollutant not in loaded_models:
            model_path = get_model_path_from_hf(pollutant)

            try:
                checkpoint = torch.load(
                    model_path,
                    map_location="cpu",
                    weights_only=True
                )
            except TypeError:
                checkpoint = torch.load(
                    model_path,
                    map_location="cpu"
                )

            model = PollutionDifferenceModel(num_classes=19, pollution_dims=1)

            if isinstance(checkpoint, dict) and "model" in checkpoint:
                model.load_state_dict(checkpoint["model"])
            elif isinstance(checkpoint, dict) and "state_dict" in checkpoint:
                model.load_state_dict(checkpoint["state_dict"])
            else:
                model.load_state_dict(checkpoint)

            model.eval()
            torch.set_grad_enabled(False)

            loaded_models[pollutant] = model
            print(f"✅ 已加载 {pollutant} 模型: {model_path}")

    return loaded_models[pollutant]


def create_request_dirs() -> Dict[str, Any]:
    request_id = uuid.uuid4().hex
    base_dir = RUNS_DIR / request_id
    input_dir = base_dir / "input"
    output_dir = base_dir / "output"
    summary_dir = base_dir / "summary"

    for d in [input_dir, output_dir, summary_dir]:
        d.mkdir(parents=True, exist_ok=True)

    return {
        "request_id": request_id,
        "base_dir": base_dir,
        "input_dir": input_dir,
        "output_dir": output_dir,
        "summary_dir": summary_dir,
    }


async def save_upload_file(upload_file: UploadFile, save_path: Path) -> None:
    content = await upload_file.read()

    if not content:
        raise HTTPException(status_code=400, detail=f"上传文件为空: {upload_file.filename}")

    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(status_code=413, detail=f"文件过大（最大 10MB）: {upload_file.filename}")

    if upload_file.content_type not in ALLOWED_CONTENT_TYPES:
        raise HTTPException(
            status_code=415,
            detail=f"不支持的文件类型 '{upload_file.content_type}'，仅支持 JPEG / PNG / WebP"
        )

    save_path.write_bytes(content)


def preprocess_image(img_np: np.ndarray) -> torch.Tensor:
    img = cv2.resize(img_np, (256, 256))
    img = img.astype(np.float32) / 255.0
    img = img.transpose(2, 0, 1)
    return torch.from_numpy(img).unsqueeze(0)


def read_rgb_image(path: Path) -> np.ndarray:
    img = cv2.imread(str(path))

    if img is None:
        raise HTTPException(status_code=400, detail=f"无法读取图像: {path.name}")

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


async def run_segmentation_async(input_dir: Path, output_dir: Path, summary_dir: Path) -> None:
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(
        None,
        partial(
            segmenter.segment,
            dir_input=str(input_dir),
            dir_image_output=str(output_dir),
            dir_summary_output=str(summary_dir)
        )
    )


def find_segmented_img(output_dir: Path, base_name: str) -> Optional[Path]:
    candidates = sorted([
        f for f in output_dir.iterdir()
        if base_name in f.name and "colored_segmented" in f.name
    ])

    return candidates[0] if candidates else None


def find_blend_img(output_dir: Path, base_name: str) -> Optional[Path]:
    candidates = sorted([
        f for f in output_dir.iterdir()
        if base_name in f.name and "blend" in f.name
    ])

    return candidates[0] if candidates else None


def copy_segmentation_outputs(output_dir: Path, request_id: str) -> Dict[str, str]:
    ref_seg_path = find_segmented_img(output_dir, "ref")
    query_seg_path = find_segmented_img(output_dir, "query")
    ref_blend_path = find_blend_img(output_dir, "ref")
    query_blend_path = find_blend_img(output_dir, "query")

    if not ref_seg_path or not query_seg_path:
        raise HTTPException(status_code=500, detail="找不到分割结果图像")

    target_ref = STATIC_RESULTS_DIR / f"{request_id}_ref_seg.png"
    target_query = STATIC_RESULTS_DIR / f"{request_id}_query_seg.png"
    target_ref_blend = STATIC_RESULTS_DIR / f"{request_id}_ref_blend.png"
    target_query_blend = STATIC_RESULTS_DIR / f"{request_id}_query_blend.png"

    shutil.copy(ref_seg_path, target_ref)
    shutil.copy(query_seg_path, target_query)

    if ref_blend_path and ref_blend_path.exists():
        shutil.copy(ref_blend_path, target_ref_blend)

    if query_blend_path and query_blend_path.exists():
        shutil.copy(query_blend_path, target_query_blend)

    return {
        "ref_seg": f"/static/results/{request_id}_ref_seg.png",
        "query_seg": f"/static/results/{request_id}_query_seg.png",
        "ref_blend": f"/static/results/{request_id}_ref_blend.png" if target_ref_blend.exists() else "",
        "query_blend": f"/static/results/{request_id}_query_blend.png" if target_query_blend.exists() else "",
    }


def infer_difference(
    model: PollutionDifferenceModel,
    ref_tensor: torch.Tensor,
    query_tensor: torch.Tensor
) -> float:
    with torch.no_grad():
        out = model(ref_tensor, query_tensor)
        return float(out.item())


def validate_ref_data(pollutant: str, ref_data: float) -> None:
    if pollutant not in POLLUTANT_RANGES:
        raise HTTPException(status_code=400, detail=f"不支持的污染物: {pollutant}")

    lo, hi = POLLUTANT_RANGES[pollutant]

    if ref_data < lo:
        raise HTTPException(status_code=422, detail=f"{pollutant} 参考值不能为负数")

    if ref_data > hi:
        raise HTTPException(
            status_code=422,
            detail=f"{pollutant} 参考值 {ref_data} 超出合理范围（最大 {hi}）"
        )


async def cleanup_old_runs(max_age_hours: int = 24) -> None:
    cutoff = datetime.now() - timedelta(hours=max_age_hours)

    if not RUNS_DIR.exists():
        return

    for run_dir in RUNS_DIR.iterdir():
        if run_dir.is_dir():
            try:
                mtime = datetime.fromtimestamp(run_dir.stat().st_mtime)
                if mtime < cutoff:
                    shutil.rmtree(run_dir, ignore_errors=True)
            except Exception:
                pass


# =========================
# 启动事件
# =========================
@app.on_event("startup")
async def startup_event():
    await cleanup_old_runs()

    try:
        for pollutant in MODEL_FILENAMES.keys():
            load_pollution_model(pollutant)

        print("✅ 所有污染预测模型预加载完成")

    except Exception as e:
        print(f"⚠️ 模型预加载失败: {e}")


# =========================
# 首页
# =========================
@app.get("/")
async def read_index():
    index_path = FRONTEND_DIR / "index.html"

    if not index_path.exists():
        raise HTTPException(status_code=500, detail="frontend/index.html not found")

    return FileResponse(str(index_path))


# =========================
# Mapillary 接口
# =========================
@app.get("/mapillary/config")
async def mapillary_config():
    return {
        "enabled": bool(MAPILLARY_ACCESS_TOKEN),
        "client_token": MAPILLARY_ACCESS_TOKEN
    }


@app.get("/mapillary/nearby")
async def mapillary_nearby(
    lat: float = Query(..., ge=-90, le=90),
    lng: float = Query(..., ge=-180, le=180),
    radius_m: int = Query(250, ge=50, le=1500),
):
    lat_delta = radius_m / 111_320.0
    lng_delta = radius_m / max(111_320.0 * math.cos(math.radians(lat)), 1.0)
    bbox = f"{lng-lng_delta},{lat-lat_delta},{lng+lng_delta},{lat+lat_delta}"

    fields = "id,computed_geometry,thumb_1024_url,thumb_2048_url,is_pano,captured_at,creator"
    payload = await mapillary_get("images", {"fields": fields, "bbox": bbox, "limit": 60})

    images = []

    for item in payload.get("data", []):
        coordinates = (item.get("computed_geometry") or {}).get("coordinates") or []

        if len(coordinates) < 2:
            continue

        distance = haversine_m(lat, lng, float(coordinates[1]), float(coordinates[0]))

        images.append({
            "id": str(item.get("id", "")),
            "lat": coordinates[1],
            "lng": coordinates[0],
            "distance_m": round(distance),
            "is_pano": bool(item.get("is_pano")),
            "captured_at": item.get("captured_at"),
            "creator": item.get("creator"),
            "thumbnail": item.get("thumb_1024_url") or item.get("thumb_2048_url"),
        })

    images.sort(key=lambda item: (not item["is_pano"], item["distance_m"]))

    return {"images": images[:12]}


@app.get("/mapillary/query-image/{image_id}")
async def mapillary_query_image(
    image_id: str,
    center_x: float = Query(0.5, ge=0, le=1),
    center_y: float = Query(0.5, ge=0, le=1),
    bearing: Optional[float] = Query(None),
):
    if not image_id.isdigit():
        raise HTTPException(status_code=400, detail="無效的 Mapillary 影像編號。")

    metadata = await mapillary_get(
        image_id,
        {
            "fields": "id,thumb_2048_url,thumb_original_url,is_pano,computed_compass_angle,captured_at,creator"
        }
    )

    image_url = metadata.get("thumb_2048_url") or metadata.get("thumb_original_url")

    if not image_url:
        raise HTTPException(status_code=404, detail="Mapillary 沒有提供這張影像。")

    try:
        async with httpx.AsyncClient(timeout=30, follow_redirects=True) as client:
            response = await client.get(image_url)

        response.raise_for_status()

        encoded = np.frombuffer(response.content, dtype=np.uint8)
        image = cv2.imdecode(encoded, cv2.IMREAD_COLOR)

        if image is None:
            raise ValueError("invalid image")

        if metadata.get("is_pano"):
            if bearing is not None and metadata.get("computed_compass_angle") is not None:
                relative_bearing = (bearing - float(metadata["computed_compass_angle"]) + 540) % 360 - 180
                center_x = (0.5 + relative_bearing / 360.0) % 1.0

            image = panorama_view(image, center_x, center_y)

        ok, output = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 92])

        if not ok:
            raise ValueError("encode failed")

    except (httpx.HTTPError, ValueError) as exc:
        raise HTTPException(status_code=502, detail="無法準備這張 Mapillary 街景。") from exc

    headers = {
        "Content-Disposition": f'inline; filename="mapillary-{image_id}.jpg"',
        "X-Mapillary-Image-Id": image_id,
    }

    return StreamingResponse(io.BytesIO(output.tobytes()), media_type="image/jpeg", headers=headers)


# =========================
# 单图预测
# =========================
@app.post("/predict")
async def predict(
    pollutant: str = Form(...),
    ref_data: float = Form(...),
    ref_img: UploadFile = File(...),
    query_img: UploadFile = File(...)
):
    try:
        validate_ref_data(pollutant, ref_data)

        paths = create_request_dirs()
        request_id = paths["request_id"]
        input_dir = paths["input_dir"]
        output_dir = paths["output_dir"]
        summary_dir = paths["summary_dir"]

        ref_path = input_dir / "ref.jpg"
        query_path = input_dir / "query.jpg"

        await save_upload_file(ref_img, ref_path)
        await save_upload_file(query_img, query_path)

        await run_segmentation_async(input_dir, output_dir, summary_dir)

        seg_urls = copy_segmentation_outputs(output_dir, request_id)

        ref_tensor = preprocess_image(read_rgb_image(ref_path))
        query_tensor = preprocess_image(read_rgb_image(query_path))

        model = load_pollution_model(pollutant)
        model_out = infer_difference(model, ref_tensor, query_tensor)

        final_pred = ref_data - model_out

        ratio_json_path = summary_dir / "pixel_ratios.json"

        if not ratio_json_path.exists():
            raise HTTPException(status_code=500, detail="分割后未找到 pixel_ratios.json")

        return {
            "status": "ok",
            "request_id": request_id,
            "pollutant": pollutant,
            "ref_data": ref_data,
            "model_out": round(model_out, 4),
            "pred_value": round(final_pred, 4),
            "ref_seg": seg_urls["ref_seg"],
            "query_seg": seg_urls["query_seg"],
            "ref_blend": seg_urls["ref_blend"],
            "query_blend": seg_urls["query_blend"],
            "ratio_json": f"/runs/{request_id}/summary/pixel_ratios.json"
        }

    except HTTPException:
        raise

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"预测失败: {str(e)}")


# =========================
# 批量预测后台任务
# =========================
async def batch_predict_task(
    request_id: str,
    pollutant: str,
    ref_data: float,
    ref_tensor: torch.Tensor,
    model: PollutionDifferenceModel,
    query_file_paths: list,
    batch_input_dir: Path
):
    results = []
    failed = []
    total = len(query_file_paths)

    with progress_lock:
        batch_progress[request_id] = {
            "total": total,
            "current": 0,
            "results": [],
            "failed": [],
            "status": "processing"
        }

    with torch.no_grad():
        for idx, file_info in enumerate(query_file_paths):
            safe_name = file_info["name"]
            query_path = Path(file_info["path"])

            try:
                query_np = read_rgb_image(query_path)
                query_tensor = preprocess_image(query_np)

                out = model(ref_tensor, query_tensor)
                model_out = float(out.item())

                final_pred = ref_data + model_out

                results.append({
                    "filename": safe_name,
                    "status": "ok",
                    "pred_value": round(final_pred, 4),
                    "model_out": round(model_out, 4),
                })

            except Exception as e:
                error_msg = f"文件：{safe_name}，错误：{str(e)}"
                print(f"【批量预测失败】{error_msg}")

                failed.append({
                    "filename": safe_name,
                    "status": "error",
                    "message": str(e)
                })

            with progress_lock:
                batch_progress[request_id]["current"] = idx + 1
                batch_progress[request_id]["results"] = results
                batch_progress[request_id]["failed"] = failed

            await asyncio.sleep(0.05)

    with progress_lock:
        batch_progress[request_id]["status"] = "completed"
        print(f"【批量任务完成】{request_id} | 成功：{len(results)} 张，失败：{len(failed)} 张")


# =========================
# 批量预测
# =========================
@app.post("/batch-predict")
async def batch_predict(
    pollutant: str = Form(...),
    ref_data: float = Form(...),
    ref_img: UploadFile = File(...),
    query_files: List[UploadFile] = File(...)
):
    try:
        validate_ref_data(pollutant, ref_data)

        if not query_files:
            raise HTTPException(status_code=400, detail="未上传任何查询图像")

        paths = create_request_dirs()
        request_id = paths["request_id"]
        batch_input_dir = paths["input_dir"]

        ref_path = batch_input_dir / "ref.jpg"
        await save_upload_file(ref_img, ref_path)

        ref_tensor = preprocess_image(read_rgb_image(ref_path))
        model = load_pollution_model(pollutant)

        query_file_paths = []

        for file in query_files:
            safe_name = os.path.basename(file.filename) if file.filename else f"{uuid.uuid4().hex}.jpg"
            query_path = batch_input_dir / safe_name

            await save_upload_file(file, query_path)

            query_file_paths.append({
                "path": str(query_path),
                "name": safe_name
            })

        asyncio.create_task(
            batch_predict_task(
                request_id=request_id,
                pollutant=pollutant,
                ref_data=ref_data,
                ref_tensor=ref_tensor,
                model=model,
                query_file_paths=query_file_paths,
                batch_input_dir=batch_input_dir
            )
        )

        return {
            "status": "processing",
            "request_id": request_id,
            "total_files": len(query_files)
        }

    except HTTPException:
        raise

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量预测启动失败: {str(e)}")


# =========================
# 批量预测进度推送
# =========================
@app.get("/progress/{request_id}")
async def get_batch_progress(request_id: str):
    async def event_generator():
        while True:
            progress = batch_progress.get(request_id, {})

            if not progress:
                yield 'data: {"error": "任务不存在"}\n\n'
                break

            progress_data = {
                "total": progress.get("total", 0),
                "current": progress.get("current", 0),
                "status": progress.get("status", "processing"),
                "results": progress.get("results", []),
                "failed": progress.get("failed", [])
            }

            yield f"data: {json.dumps(progress_data, ensure_ascii=False)}\n\n"

            if progress.get("status") in ["completed", "failed"]:
                break

            await asyncio.sleep(0.1)

    return StreamingResponse(event_generator(), media_type="text/event-stream")


# =========================
# 健康检查
# =========================
@app.get("/health")
async def health_check():
    return {"status": "ok"}


# =========================
# 启动
# =========================
if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)