from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi import Request
import io
from PIL import Image
import torch
from torchvision import transforms, models
import torch.nn as nn
import os
import math
import uvicorn

app = FastAPI(title="Demographic Analysis API")

app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

RACE_LABELS = ['East Asian', 'Indian', 'Black', 'White', 'Middle Eastern', 'Latino_Hispanic', 'Southeast Asian']

class VGG16Multi(nn.Module):
    def __init__(self):
        super().__init__()
        vgg = models.vgg16(weights=None)
        self.features = vgg.features
        self.avgpool = vgg.avgpool
        self.flatten = nn.Flatten()
        self.fc = nn.Sequential(
            nn.Linear(512*7*7, 512), nn.ReLU(), nn.Dropout(0.5)
        )
        self.gender = nn.Linear(512, 1)          # raw logit — sigmoid applied at inference
        self.age    = nn.Sequential(nn.Linear(512, 1), nn.Sigmoid())
        self.race   = nn.Linear(512, 7)

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = self.flatten(x)
        x = self.fc(x)
        gender = self.gender(x)
        age = self.age(x)
        race = self.race(x)
        return gender, age, race

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = VGG16Multi()

model_path = os.path.join("train", "model.pth")
if os.path.exists(model_path):
    model.load_state_dict(torch.load(model_path, map_location=device))
    print(f"Loaded model weights from {model_path}")
else:
    print(f"Warning: Model weights not found at {model_path}. Using random weights.")

model = model.to(device)
model.eval()

# Load temperature scaling for gender if available
GENDER_TEMP = 1.0
temp_path = os.path.join('train','gender_temp.txt')
if os.path.exists(temp_path):
    try:
        with open(temp_path,'r') as f:
            GENDER_TEMP = float(f.read().strip())
            print(f'Using gender temperature: {GENDER_TEMP}')
    except Exception:
        print('Failed to read gender_temp.txt, using T=1.0')

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.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
    """Serve the frontend HTML."""
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """Handle image upload and run inference."""
    try:
        if not file.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="File provided is not an image.")

        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert('RGB')
        input_tensor = transform(image).unsqueeze(0).to(device)
        
        with torch.no_grad():
            out_gender, out_age, out_race = model(input_tensor)
            # Apply sigmoid to raw logit from gender head
            gender_prob = float(torch.sigmoid(out_gender.squeeze().cpu()).item())
            # apply temperature scaling using stored temperature
            eps = 1e-6
            p = min(max(gender_prob, eps), 1 - eps)
            logit = math.log(p / (1.0 - p))
            # avoid division by zero
            T = GENDER_TEMP if GENDER_TEMP > 0 else 1.0
            calibrated_p = 1.0 / (1.0 + math.exp(-logit / T))
            female_prob = float(calibrated_p)
            male_prob = float(1.0 - calibrated_p)
            predicted_gender = "Female" if female_prob > 0.5 else "Male"

            # Age (regression output scaled back to years)
            raw_age = float(out_age.squeeze().cpu().item())
            predicted_age = raw_age * 75.0

            # Race probabilities (softmax over logits)
            race_probs = torch.softmax(out_race, dim=1)[0].cpu().tolist()
            race_idx = int(torch.tensor(race_probs).argmax().item())
            predicted_race = RACE_LABELS[race_idx]

            # Prepare race probabilities mapping
            race_probs_map = {RACE_LABELS[i]: round(race_probs[i], 3) for i in range(len(RACE_LABELS))}

            return {
                "gender": predicted_gender,
                "gender_confidence": round(max(female_prob, male_prob), 3),
                "gender_probs": {"Female": round(female_prob, 3), "Male": round(male_prob, 3)},
                "age": round(predicted_age, 1),
                "race": predicted_race,
                "race_probs": race_probs_map
            }
            
    except Exception as e:
        print(f"Inference error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)