from fastapi import FastAPI,UploadFile, File
from typing import Annotated
from fastapi.responses import JSONResponse
import uvicorn
from pydantic import BaseModel

import requests
import os 
import pandas as pd 

app = FastAPI()

import unicodedata
import string

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms 
from PIL import Image
import torchvision.transforms as transforms

import io

all_letters = string.ascii_letters + " .,;'"
n_letters = len(all_letters)

all_categories = ['Arabic','Chinese','Czech','Dutch','English','French','German','Greek',
                'Irish','Italian','Japanese','Korean','Polish','Portuguese','Russian','Scottish',
                'Spanish','Vietnamese']

# Turn a Unicode string to plain ASCII, thanks to https://stackoverflow.com/a/518232/2809427
def unicodeToAscii(s):
    return ''.join(
        c for c in unicodedata.normalize('NFD', s)
        if unicodedata.category(c) != 'Mn'
        and c in all_letters
    )

# Find letter index from all_letters, e.g. "a" = 0
def letterToIndex(letter):
    return all_letters.find(letter)


# Turn a line into a <line_length x 1 x n_letters>,
# or an array of one-hot letter vectors
def lineToTensor(line):
    tensor = torch.zeros(len(line), 1, n_letters)
    for li, letter in enumerate(line):
        tensor[li][0][letterToIndex(letter)] = 1
    return tensor

# Just return an output given a line
def evaluate_model(line_tensor):
    model = torch.jit.load("torchscript_classify_names_lstm.pt")
    
    hidden = (torch.zeros(1, 1, 128),torch.zeros(1, 1, 128))
    output = model(line_tensor,hidden)

    return output

def classify_lastname(last_name):
    # Converting to Ascii and capitalizing first letter
    last_name = unicodeToAscii(last_name)
    last_name = last_name.title()

    # Converting name to tensor
    line_tensor = lineToTensor(last_name)
    output = evaluate_model(line_tensor)

    # Grabbing top probability and category
    top_prob, top_cat = torch.topk(output,1)
    prob = torch.exp(top_prob[0])
    cat = top_cat[0]

    model_output = {}
    model_output[all_categories[cat[0].item()]] = round(prob[0].item(),2)

    return model_output
    
# Define a request model
class Item(BaseModel):
    text: str

# Endpoint to return all products by profit
@app.get("/products-by-profit")
def get_products_by_profit(top: int = 10, ascending: bool = False):
    df = pd.read_csv("Sample-Superstore.csv",encoding='ISO-8859-1')

    sorted_products = (
        df.groupby("Product Name")["Profit"]
        .sum()
        .sort_values(ascending=ascending)
        .head(top)
        .reset_index()
        .rename(columns={"Product Name": "product", "Profit": "profit"})
    )
    return sorted_products.to_dict(orient="records")

@app.post("/classifyname")
async def classify_name(lastname: Item):
    return JSONResponse(content=classify_lastname(lastname.text), status_code=201)
    
def transform_img(img):
    # Transformations that will be applied
    the_transform = transforms.Compose([
            transforms.Resize((224,224)),
            transforms.CenterCrop((224,224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
    ])
    return the_transform(img)

    # Returns string with class and probability
def classify_img(img):
    class_names = ['AIR COMPRESSOR',
 'ALTERNATOR',
 'BATTERY',
 'BRAKE CALIPER',
 'BRAKE PAD',
 'BRAKE ROTOR',
 'CAMSHAFT',
 'CARBERATOR',
 'CLUTCH PLATE',
 'COIL SPRING',
 'CRANKSHAFT',
 'CYLINDER HEAD',
 'DISTRIBUTOR',
 'ENGINE BLOCK',
 'ENGINE VALVE',
 'FUEL INJECTOR',
 'FUSE BOX',
 'GAS CAP',
 'HEADLIGHTS',
 'IDLER ARM',
 'IGNITION COIL',
 'INSTRUMENT CLUSTER',
 'LEAF SPRING',
 'LOWER CONTROL ARM',
 'MUFFLER',
 'OIL FILTER',
 'OIL PAN',
 'OIL PRESSURE SENSOR',
 'OVERFLOW TANK',
 'OXYGEN SENSOR',
 'PISTON',
 'PRESSURE PLATE',
 'RADIATOR',
 'RADIATOR FAN',
 'RADIATOR HOSE',
 'RADIO',
 'RIM',
 'SHIFT KNOB',
 'SIDE MIRROR',
 'SPARK PLUG',
 'SPOILER',
 'STARTER',
 'TAILLIGHTS',
 'THERMOSTAT',
 'TORQUE CONVERTER',
 'TRANSMISSION',
 'VACUUM BRAKE BOOSTER',
 'VALVE LIFTER',
 'WATER PUMP',
 'WINDOW REGULATOR']
    model = torch.jit.load("car_part_traced_classifier_resnet50.ptl")
    # Applying transformation to the image
    model_img = transform_img(img)
    model_img = model_img.view(1,3,224,224)
    
    # Running image through the model
    model.eval()
    with torch.no_grad():
        result = model(model_img)
    
    # Converting values to softmax values 
    result = F.softmax(result,dim=1)
    # Grabbing top 3 indices and probabilities for each index
    top3_prob, top3_catid = torch.topk(result,3)
    
    # Dictionary I will display
    model_output = {}
    for i in range(top3_prob.size(1)):
        
        model_output[class_names[top3_catid[0][i].item()]] = top3_prob[0][i].item()
    print(model_output)
    return model_output 

    
@app.post("/classifycarpart")
async def upload_image(file: UploadFile = File(...)):
    try:
        # Read and convert the image to a PIL Image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")
     
        return JSONResponse(content=classify_img(image), status_code=201)

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

@app.get("/")
def api_home():
    return JSONResponse({'detail': 'Welcome to FastAPI!'}, status_code=200)