from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import pickle

app = FastAPI(title="Transaction Classifier API", description="API for classifying banking transactions.")
print("FastAPI app initialized...")

max_len = 20

class HierarchicalPrediction(tf.keras.layers.Layer):
    def __init__(self, num_subcategories, cat_to_subcat_tensor, max_subcats_per_cat, **kwargs):
        super(HierarchicalPrediction, self).__init__(**kwargs)
        self.num_subcategories = num_subcategories
        self.cat_to_subcat_tensor = cat_to_subcat_tensor
        self.max_subcats_per_cat = max_subcats_per_cat
        self.subcategory_dense = Dense(num_subcategories, activation=None)

    def build(self, input_shape):
        super(HierarchicalPrediction, self).build(input_shape)

    def call(self, inputs):
        lstm_output, category_probs = inputs
        subcat_logits = self.subcategory_dense(lstm_output)
        batch_size = tf.shape(category_probs)[0]
        predicted_categories = tf.argmax(category_probs, axis=1)
        valid_subcat_indices = tf.gather(self.cat_to_subcat_tensor, predicted_categories)
        batch_indices = tf.range(batch_size)
        batch_indices_expanded = tf.tile(batch_indices[:, tf.newaxis], [1, self.max_subcats_per_cat])
        update_indices = tf.stack([batch_indices_expanded, valid_subcat_indices], axis=-1)
        update_indices = tf.reshape(update_indices, [-1, 2])
        valid_mask = tf.not_equal(valid_subcat_indices, -1)
        valid_indices = tf.boolean_mask(update_indices, tf.reshape(valid_mask, [-1]))
        updates = tf.ones(tf.shape(valid_indices)[0], dtype=tf.float32)
        mask = tf.scatter_nd(valid_indices, updates, [batch_size, self.num_subcategories])
        masked_logits = subcat_logits * mask + (1 - mask) * tf.float32.min
        return tf.nn.softmax(masked_logits)

    def get_config(self):
        config = super(HierarchicalPrediction, self).get_config()
        config.update({
            'num_subcategories': self.num_subcategories,
            'cat_to_subcat_tensor': self.cat_to_subcat_tensor.numpy(),
            'max_subcats_per_cat': self.max_subcats_per_cat
        })
        return config

    @classmethod
    def from_config(cls, config):
        config['cat_to_subcat_tensor'] = tf.constant(config['cat_to_subcat_tensor'], dtype=tf.int32)
        return cls(**config)

tf.keras.utils.get_custom_objects()['HierarchicalPrediction'] = HierarchicalPrediction

def load_resources():
    print("Loading BiLSTM model...")
    model = tf.keras.models.load_model('model.h5', custom_objects={'HierarchicalPrediction': HierarchicalPrediction})
    print("Loading tokenizer...")
    with open('tokenizer.pkl', 'rb') as f:
        tokenizer = pickle.load(f)
    print("Loading category label encoder...")
    with open('le_category.pkl', 'rb') as f:
        le_category = pickle.load(f)
    print("Loading subcategory label encoder...")
    with open('le_subcategory.pkl', 'rb') as f:
        le_subcategory = pickle.load(f)
    print(f"Num categories: {len(le_category.classes_)}, Num subcategories: {len(le_subcategory.classes_)}")
    return model, tokenizer, le_category, le_subcategory

class TransactionRequest(BaseModel):
    description: str

class PredictionResponse(BaseModel):
    category: str
    subcategory: str
    category_confidence: float
    subcategory_confidence: float

@app.get("/")
async def root():
    return {"message": "Welcome to the Transaction Classifier API. Use POST /predict to classify transactions."}

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: TransactionRequest):
    try:
        model, tokenizer, le_category, le_subcategory = load_resources()
        
        seq = tokenizer.texts_to_sequences([request.description])
        pad = pad_sequences(seq, maxlen=max_len, padding='post')
        pred = model.predict(pad, verbose=0)
        
        cat_probs = pred[0][0]
        subcat_probs = pred[1][0]
        cat_idx = np.argmax(cat_probs)
        subcat_idx = np.argmax(subcat_probs)
        
        cat_pred = le_category.inverse_transform([cat_idx])[0]
        subcat_pred = le_subcategory.inverse_transform([subcat_idx])[0]
        cat_conf = float(cat_probs[cat_idx] * 100)
        subcat_conf = float(subcat_probs[subcat_idx] * 100)
        
        print(f"Predicted: category={cat_pred}, subcategory={subcat_pred}")
        return {
            "category": cat_pred,
            "subcategory": subcat_pred,
            "category_confidence": cat_conf,
            "subcategory_confidence": subcat_conf
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

print("API ready with lazy-loaded resources.")