# =========================================
# IMPORTS
# =========================================
import os
import cv2
import nibabel as nib
import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models

from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

# =========================================
# CONFIG
# =========================================
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

IMG_SIZE = 128
NUM_CLASSES = 3
NUM_SLICES = 32

LABELS = [
    "Control",
    "Prodromal",
    "Parkinsons"
]

print("Using device:", DEVICE)

# =========================================
# FASTAPI
# =========================================
app = FastAPI(
    title="Parkinson DATSCAN Ensemble API",
    version="1.0"
)

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

# =========================================
# LOAD NIFTI
# =========================================
def load_nifti_from_bytes(file_bytes):

    temp_path = "temp_upload.nii"

    with open(temp_path, "wb") as f:
        f.write(file_bytes)

    volume = nib.load(temp_path).get_fdata()

    volume = np.squeeze(volume)

    return volume

# =========================================
# PREPROCESS 2D
# =========================================
def preprocess_2d(volume):

    depth = volume.shape[2]

    idx1 = np.linspace(
        0,
        depth // 3 - 1,
        10
    ).astype(int)

    idx2 = np.linspace(
        depth // 3,
        2 * depth // 3 - 1,
        10
    ).astype(int)

    idx3 = np.linspace(
        2 * depth // 3,
        depth - 1,
        12
    ).astype(int)

    def make_channel(indices):

        slices = []

        for i in indices:

            sl = volume[:, :, i]

            sl = sl - sl.min()
            sl = sl / (sl.max() + 1e-6)

            sl = cv2.resize(
                sl,
                (IMG_SIZE, IMG_SIZE)
            )

            slices.append(sl)

        return np.mean(slices, axis=0)

    r = make_channel(idx1)
    g = make_channel(idx2)
    b = make_channel(idx3)

    img = np.stack([r, g, b], axis=0)

    tensor = torch.tensor(
        img,
        dtype=torch.float32
    )

    tensor = tensor.unsqueeze(0)

    return tensor

# =========================================
# PREPROCESS 3D
# =========================================
def preprocess_3d(volume):

    depth = volume.shape[2]

    indices = np.linspace(
        0,
        depth - 1,
        NUM_SLICES
    ).astype(int)

    slices = []

    for i in indices:

        sl = volume[:, :, i]

        sl = sl - sl.min()
        sl = sl / (sl.max() + 1e-6)

        sl = cv2.resize(
            sl,
            (IMG_SIZE, IMG_SIZE)
        )

        slices.append(sl)

    vol = np.stack(slices)

    tensor = torch.tensor(
        vol,
        dtype=torch.float32
    )

    tensor = tensor.unsqueeze(0)
    tensor = tensor.unsqueeze(0)

    return tensor

# =========================================
# DENSENET121
# =========================================
class DenseNet121Model(nn.Module):

    def __init__(self):

        super().__init__()

        self.base = models.densenet121(
            weights=None
        )

        self.base.classifier = nn.Linear(
            self.base.classifier.in_features,
            NUM_CLASSES
        )

    def forward(self, x):

        return self.base(x)

# =========================================
# DENSENET169
# =========================================
class DenseNet169Model(nn.Module):

    def __init__(self):

        super().__init__()

        self.base = models.densenet169(
            weights=None
        )

        self.base.classifier = nn.Linear(
            self.base.classifier.in_features,
            NUM_CLASSES
        )

    def forward(self, x):

        return self.base(x)

# =========================================
# DENSENET201
# =========================================
class DenseNet201Model(nn.Module):

    def __init__(self):

        super().__init__()

        self.base = models.densenet201(
            weights=None
        )

        self.base.classifier = nn.Linear(
            self.base.classifier.in_features,
            NUM_CLASSES
        )

    def forward(self, x):

        return self.base(x)

# =========================================
# 3D CNN
# =========================================
class CNN3D(nn.Module):

    def __init__(self):

        super().__init__()

        self.net = nn.Sequential(

            nn.Conv3d(
                1,
                16,
                kernel_size=3,
                padding=1
            ),

            nn.ReLU(),

            nn.MaxPool3d(2),

            nn.Conv3d(
                16,
                32,
                kernel_size=3,
                padding=1
            ),

            nn.ReLU(),

            nn.MaxPool3d(2),

            nn.Conv3d(
                32,
                64,
                kernel_size=3,
                padding=1
            ),

            nn.ReLU(),

            nn.MaxPool3d(2)
        )

        self.fc = nn.Sequential(

            nn.Linear(
                64 * 4 * 16 * 16,
                256
            ),

            nn.ReLU(),

            nn.Dropout(0.3),

            nn.Linear(
                256,
                NUM_CLASSES
            )
        )

    def forward(self, x):

        x = self.net(x)

        x = x.view(
            x.size(0),
            -1
        )

        x = self.fc(x)

        return x

# =========================================
# LOAD MODELS
# =========================================
def load_model(model, path):

    print(f"Loading {path}")

    state = torch.load(
        path,
        map_location=DEVICE
    )

    model.load_state_dict(
        state,
        strict=False
    )

    model.to(DEVICE)

    model.eval()

    print(f"Loaded {path}")

    return model

# =========================================
# LOAD ALL
# =========================================
model121 = load_model(
    DenseNet121Model(),
    "densenet121.pth"
)

model169 = load_model(
    DenseNet169Model(),
    "densenet169.pth"
)

model201 = load_model(
    DenseNet201Model(),
    "densenet201.pth"
)

model3d = load_model(
    CNN3D(),
    "cnn3d.pth"
)

# =========================================
# SINGLE PREDICTION
# NOTE: All confidence and probability values are 0.0–1.0 (NOT percentages)
# =========================================
def predict_model(model, tensor):

    tensor = tensor.to(DEVICE)

    with torch.no_grad():

        out = model(tensor)

        probs = torch.softmax(
            out,
            dim=1
        )[0]

    pred_idx = torch.argmax(
        probs
    ).item()

    result = {

        "prediction":
            LABELS[pred_idx],

        "class_id":
            pred_idx,

        # Confidence as 0.0–1.0 decimal (NOT percentage)
        "confidence":
            round(
                float(
                    probs[pred_idx]
                ),
                4
            ),

        "probabilities": {

            LABELS[i]:
                round(
                    float(
                        probs[i]
                    ),
                    4
                )

            for i in range(NUM_CLASSES)
        }
    }

    return result, probs

# =========================================
# HEALTH ENDPOINT
# =========================================
@app.get("/health")
def health():

    return {
        "status": "running",
        "service": "Parkinson DaTscan Ensemble API",
        "models_loaded": {
            "densenet121": True,
            "densenet169": True,
            "densenet201": True,
            "cnn3d": True,
        },
        "device": str(DEVICE),
        "classes": LABELS,
    }

# =========================================
# ROOT
# =========================================
@app.get("/")
def home():

    return {

        "message":
            "Parkinson DATSCAN Ensemble API Running",

        "classes":
            LABELS,

        "endpoints": [
            "/health",
            "/predict/densenet121",
            "/predict/densenet169",
            "/predict/densenet201",
            "/predict/cnn3d",
            "/predict/ensemble",
        ]
    }

# =========================================
# DENSENET121 ENDPOINT
# =========================================
@app.post("/predict/densenet121")
async def predict_121(
    file: UploadFile = File(...)
):

    file_bytes = await file.read()

    volume = load_nifti_from_bytes(
        file_bytes
    )

    tensor = preprocess_2d(volume)

    result, _ = predict_model(
        model121,
        tensor
    )

    return JSONResponse(result)

# =========================================
# DENSENET169 ENDPOINT
# =========================================
@app.post("/predict/densenet169")
async def predict_169(
    file: UploadFile = File(...)
):

    file_bytes = await file.read()

    volume = load_nifti_from_bytes(
        file_bytes
    )

    tensor = preprocess_2d(volume)

    result, _ = predict_model(
        model169,
        tensor
    )

    return JSONResponse(result)

# =========================================
# DENSENET201 ENDPOINT
# =========================================
@app.post("/predict/densenet201")
async def predict_201(
    file: UploadFile = File(...)
):

    file_bytes = await file.read()

    volume = load_nifti_from_bytes(
        file_bytes
    )

    tensor = preprocess_2d(volume)

    result, _ = predict_model(
        model201,
        tensor
    )

    return JSONResponse(result)

# =========================================
# 3D CNN ENDPOINT
# =========================================
@app.post("/predict/cnn3d")
async def predict_cnn3d(
    file: UploadFile = File(...)
):

    file_bytes = await file.read()

    volume = load_nifti_from_bytes(
        file_bytes
    )

    tensor = preprocess_3d(volume)

    result, _ = predict_model(
        model3d,
        tensor
    )

    return JSONResponse(result)

# =========================================
# ENSEMBLE ENDPOINT
# Returns ensemble confidence as 0.0-1.0 with individual model breakdown
# =========================================
@app.post("/predict/ensemble")
async def predict_ensemble(
    file: UploadFile = File(...)
):

    file_bytes = await file.read()

    volume = load_nifti_from_bytes(
        file_bytes
    )

    tensor2d = preprocess_2d(volume)

    tensor3d = preprocess_3d(volume)

    r121, p121 = predict_model(
        model121,
        tensor2d
    )

    r169, p169 = predict_model(
        model169,
        tensor2d
    )

    r201, p201 = predict_model(
        model201,
        tensor2d
    )

    r3d, p3d = predict_model(
        model3d,
        tensor3d
    )

    avg_probs = (
        p121 +
        p169 +
        p201 +
        p3d
    ) / 4

    pred_idx = torch.argmax(
        avg_probs
    ).item()

    # All probability values as 0.0–1.0
    final_result = {

        "prediction":
            LABELS[pred_idx],

        # Ensemble confidence as 0.0–1.0 (NOT percentage)
        "confidence":
            round(
                float(
                    avg_probs[pred_idx]
                ),
                4
            ),

        "probabilities": {

            LABELS[i]:
                round(
                    float(
                        avg_probs[i]
                    ),
                    4
                )

            for i in range(NUM_CLASSES)
        },

        "individual_models": {

            "DenseNet121":
                r121,

            "DenseNet169":
                r169,

            "DenseNet201":
                r201,

            "CNN3D":
                r3d
        }
    }

    return JSONResponse(final_result)