import os
import io
import torch
import torch.nn.functional as F
from flask import Flask, render_template, request, jsonify
from torchvision import transforms
from PIL import Image

from convnext_mlp import ConvNextMLP
from convnext_kan import ConvNextKAN

app = Flask(__name__)

CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

inference_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

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

def load_model_weights(model, path):
    if os.path.exists(path):
        model.load_state_dict(torch.load(path, map_location=device))
        print(f"Pth file of {path} successfully loaded!")
    else:
        print(f"Pth file of {path} not found. Please check the path.")
    
    model.eval()
    return model.to(device)

FREEZE_BACKBONE = True
UNFREEZE_LAST_STAGE = False

model_mlp = ConvNextMLP(
    num_classes=10, 
    head_depth=3, 
    hidden_dim_1=512, 
    hidden_dim_2=512, 
    head_style="rakyan",
    freeze_backbone=FREEZE_BACKBONE,
    unfreeze_last_stage=UNFREEZE_LAST_STAGE
)
model_mlp = load_model_weights(model_mlp, "weights/convnext_mlp_cifar10.pth")

model_kan = ConvNextKAN(
    num_classes=10, 
    head_depth=3, 
    hidden_dim_1=512, 
    hidden_dim_2=512, 
    head_style="rakyan",
    freeze_backbone=FREEZE_BACKBONE,
    unfreeze_last_stage=UNFREEZE_LAST_STAGE
)
model_kan = load_model_weights(model_kan, "weights/convnext_kan_cifar10.pth")

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/api/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"status": "error", "message": "No file"}), 400
    
    file = request.files['file']
    img_bytes = file.read()
    
    try:
        image = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        tensor_img = inference_transform(image).unsqueeze(0).to(device)

        with torch.no_grad():
            out_1 = model_mlp(tensor_img)
            prob_1 = F.softmax(out_1, dim=1)
            conf_1, idx_1 = torch.max(prob_1, 1)

            out_2 = model_kan(tensor_img)
            prob_2 = F.softmax(out_2, dim=1)
            conf_2, idx_2 = torch.max(prob_2, 1)

        p1_list = prob_1[0].tolist()
        p2_list = prob_2[0].tolist()

        all_p1 = [{"class": CLASSES[i], "confidence": round(p1_list[i] * 100, 2)} for i in range(10)]
        all_p2 = [{"class": CLASSES[i], "confidence": round(p2_list[i] * 100, 2)} for i in range(10)]
        
        all_p1.sort(key=lambda x: x['confidence'], reverse=True)
        all_p2.sort(key=lambda x: x['confidence'], reverse=True)

        is_outlier = False
        if conf_1.item() < 0.35 and conf_2.item() < 0.35:
            is_outlier = True

        return jsonify({
            "status": "success",
            "is_outlier": is_outlier,
            "model_1": {
                "class": CLASSES[idx_1.item()], 
                "confidence": round(conf_1.item() * 100, 2), 
                "all_probs": all_p1
            },
            "model_2": {
                "class": CLASSES[idx_2.item()], 
                "confidence": round(conf_2.item() * 100, 2), 
                "all_probs": all_p2
            }
        })
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)