import datetime
import threading
import numpy as np
import pandas as pd
from io import BytesIO
from fastapi import FastAPI, HTTPException
from contextlib import asynccontextmanager
from pydantic import BaseModel
from transformers import AutoTokenizer
from transformers import get_linear_schedule_with_warmup, AutoTokenizer
from sklearn.model_selection import train_test_split
from gensim.utils import simple_preprocess
from sklearn.model_selection import StratifiedKFold
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from typing import List 
from minio.error import S3Error

from bert_model import BERTClassifier
from utils import seed_everything
from minio_client import minio_client, download_latest_model
from constants import PHOBERTBASE_DIR, CLASSES, DEVICE, MODEL_CHECKPOINT, EPOCH, \
                        MAX_TOKEN_LENGTH, BATCH_SIZE, LOAD_DATA_WORKER, MINIO_MODEL_TRAINED, \
                        MINIO_DATA_LABELED, TEST_RATIO, K_FOLD
from logger import logger
from preprocess import preprocess_row

# Global variables to manage the training thread and the stop flag
train_thread = None
is_training = False
stop_training_flag = threading.Event()

def release_training_vram():
    # Giải phóng bộ nhớ GPU của quá trình training
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

class ClassifyDataset(Dataset):
    def __init__(self, df, tokenizer, classes, max_len):
        self.df = df
        self.max_len = max_len
        self.tokenizer = tokenizer
        self.classes = classes
    
    def __len__(self):
        return len(self.df)

    def __getitem__(self, index):
        """
        To customize dataset, inherit from Dataset class and implement
        __len__ & __getitem__
        __getitem__ should return 
            data:
                input_ids
                attention_masks
                text
                targets
        """
        row = self.df.iloc[index]
        text, label = self.get_input_data(row)

        # Encode_plus will:
        # (1) split text into token
        # (2) Add the '[CLS]' and '[SEP]' token to the start and end
        # (3) Truncate/Pad sentence to max length
        # (4) Map token to their IDS
        # (5) Create attention mask
        # (6) Return a dictionary of outputs
        encoding = self.tokenizer.encode_plus(
            text,
            truncation=True,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            return_attention_mask=True,
            return_token_type_ids=False,
            return_tensors='pt',
        )
        
        return {
            'text': text,
            'input_ids': encoding['input_ids'].flatten(),
            'attention_masks': encoding['attention_mask'].flatten(),
            'targets': torch.tensor(label, dtype=torch.long),
        }


    def labelencoder(self,text):
        for i in range(len(self.classes)):
            if text == self.classes[i]:
                return i

    def get_input_data(self, row):
        # Preprocessing: {remove icon, special character, lower}
        text = row['text']
        text = ' '.join(simple_preprocess(text))
        label = self.labelencoder(row['label'])

        return text, label

def eval(model, valid_loader, criterion, device):
    model.eval()
    losses = []
    correct = 0

    with torch.no_grad():
        for data in valid_loader:
            input_ids = data['input_ids'].to(device)
            attention_mask = data['attention_masks'].to(device)
            targets = data['targets'].to(device)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask
            )

            _, pred = torch.max(outputs, dim=1)

            loss = criterion(outputs, targets)
            correct += torch.sum(pred == targets)
            losses.append(loss.item())
    
    logger.info(f'Valid Accuracy: {correct.double()/len(valid_loader.dataset)} Loss: {np.mean(losses)}')
    return correct.double()/len(valid_loader.dataset)

def prepare_loaders(df, fold, tokenizer, classes, max_len, batch_size, num_workers):
    df_train = df[df.kfold != fold].reset_index(drop=True)
    df_valid = df[df.kfold == fold].reset_index(drop=True)
    
    train_dataset = ClassifyDataset(df_train, tokenizer, classes, max_len=max_len)
    valid_dataset = ClassifyDataset(df_valid, tokenizer, classes, max_len=max_len)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    
    return train_loader, valid_loader


def train_model(skf, train_df, tokenizer, use_pretrain=False):
    global stop_training_flag, is_training
    train_version = datetime.datetime.now().strftime("%d-%m-%Y_%Hh%Mm")
    
    writer = SummaryWriter(log_dir=f'/runs/{train_version}')

    logger.info("START TRAINING PROCESS")
    model = BERTClassifier(model_bert=PHOBERTBASE_DIR, n_classes=len(CLASSES)).to(DEVICE)
    if use_pretrain == True:
        model.load_state_dict(torch.load(MODEL_CHECKPOINT))

    epoch_each_fold = int(EPOCH/skf.n_splits)
    cur_epoch = 0

    for fold in range(skf.n_splits):
        if stop_training_flag.is_set():
            logger.info("Training stopped.")
            writer.close()
            return "Training stopped"
        logger.info(f'-----------Fold: {fold+1} ------------------')
        train_loader, valid_loader = prepare_loaders(train_df, fold=fold, tokenizer=tokenizer, classes=CLASSES, \
                                                     max_len=MAX_TOKEN_LENGTH, batch_size=BATCH_SIZE, num_workers=LOAD_DATA_WORKER)
        criterion = nn.CrossEntropyLoss()
        # Recommendation by BERT: lr: 5e-5, 2e-5, 3e-5
        # Batchsize: 16, 32
        optimizer = AdamW(model.parameters(), lr=2e-5)
        
        lr_scheduler = get_linear_schedule_with_warmup( 
                    optimizer, 
                    num_warmup_steps=0, 
                    num_training_steps=len(train_loader)*EPOCH
                )
        best_acc = 0
        for e in range(epoch_each_fold):
            if stop_training_flag.is_set():
                logger.info("Training stopped.")
                writer.close()
                return "Training stopped"
            logger.info(f'Fold {fold+1} Epoch {cur_epoch+1}/{EPOCH}')
            logger.info('-'*30)
            # Train ----------------------------------------------------------------------------------
            model.train()
            losses = []
            correct = 0

            for data in train_loader:
                if stop_training_flag.is_set():
                    logger.info("Training stopped.")
                    writer.close()
                    return "Training stopped"
                if data is not None:         
                    input_ids = data['input_ids'].to(DEVICE)
                    attention_mask = data['attention_masks'].to(DEVICE)
                    targets = data['targets'].to(DEVICE)

                    optimizer.zero_grad()
                    outputs = model(
                        input_ids=input_ids,
                        attention_mask=attention_mask
                    )

                    loss = criterion(outputs, targets)
                    _, pred = torch.max(outputs, dim=1)

                    correct += torch.sum(pred == targets)
                    losses.append(loss.item())
                    loss.backward()
                    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                    optimizer.step()
                    lr_scheduler.step()
                else:
                    logger.warning("Warning: Empty data received from data loader. Skipping this iteration.")

            train_acc = correct.double()/len(train_loader.dataset)
            train_loss = np.mean(losses)
            logger.info(f'Train Accuracy: {train_acc} Loss: {train_loss}')
            writer.add_scalar("Accuracy/train", train_acc, cur_epoch)
            writer.add_scalar("Loss/train", train_loss, cur_epoch)
            # End train ----------------------------------------------------------------------------------
            if stop_training_flag.is_set():
                logger.info("Training stopped.")
                writer.close()
                return "Training stopped"
            # Valid --------------------------------------------------------------------------------------
            model.eval()
            losses = []
            correct = 0

            with torch.no_grad():
                for data in valid_loader:
                    if stop_training_flag.is_set():
                        logger.info("Training stopped.")
                        writer.close()
                        return "Training stopped"
                    input_ids = data['input_ids'].to(DEVICE)
                    attention_mask = data['attention_masks'].to(DEVICE)
                    targets = data['targets'].to(DEVICE)

                    outputs = model(
                        input_ids=input_ids,
                        attention_mask=attention_mask
                    )

                    _, pred = torch.max(outputs, dim=1)

                    loss = criterion(outputs, targets)
                    correct += torch.sum(pred == targets)
                    losses.append(loss.item())
            
            val_acc = correct.double()/len(valid_loader.dataset)
            val_loss = np.mean(losses)
            logger.info(f'Valid Accuracy: {val_acc} Loss: {val_loss}')
            writer.add_scalar("Accuracy/valid", val_acc, cur_epoch)
            writer.add_scalar("Loss/valid", val_loss, cur_epoch)
            # End valid ----------------------------------------------------------------------------------
            
            if stop_training_flag.is_set():
                logger.info("Training stopped.")
                writer.close()
                return "Training stopped"
            
            # Save checkpoint
            torch.save(model.state_dict(), f'/src/phobert_last.pth')

            if val_acc > best_acc:
                torch.save(model.state_dict(), f'/src/phobert_best.pth')
                best_acc = val_acc
                checkpoint_best = f'phobert_best_{train_version}.pth'
                try:
                    upload_best = minio_client.fput_object(
                        MINIO_MODEL_TRAINED,      # Bucket name
                        checkpoint_best,      # Object name (name in MinIO)
                        '/src/phobert_best.pth'      # Path to the file you want to upload
                    )
                except S3Error as exc:
                    logger.error(f"Error occurred: {exc}")
                logger.info(f"File uploaded successfully. {checkpoint_best}")

            if stop_training_flag.is_set():
                logger.info("Training stopped.")
                writer.close()
                return "Training stopped"
            
            # Upload file
            checkpoint_last = f'phobert_last_{train_version}.pth'
            try:
                upload_last = minio_client.fput_object(
                    MINIO_MODEL_TRAINED,      # Bucket name
                    checkpoint_last,      # Object name (name in MinIO)
                    '/src/phobert_last.pth'      # Path to the file you want to upload
                )
            except S3Error as exc:
                logger.error(f"Error occurred: {exc}")
            logger.info(f"File uploaded successfully. {checkpoint_last}")

            cur_epoch = cur_epoch + 1

    logger.info("Training completed.")
    writer.close()
    release_training_vram()
    is_training = False
    return "Training completed"


    
class TrainingRequest(BaseModel):
    pretrain: str

class TrainingResponse(BaseModel):
    status: str
    message: str

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(f'CUDA available: {torch.cuda.is_available()}')
    yield
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

app = FastAPI(lifespan=lifespan)

@app.post("/start-training", response_model=TrainingResponse)
async def start_training(request: TrainingRequest):
    global train_thread, stop_training_flag, is_training
    if is_training == True:
        raise HTTPException(status_code=400, detail="Training already in progress")
    is_training = True
    
    use_pretrain = False

    if request.pretrain == "":
        logger.info("Training from zero")
    elif request.pretrain == "latest":
        try:
            latest_model = download_latest_model()
        except Exception as exc:
            is_training = False
            raise HTTPException(status_code=500, detail=str(exc))
        use_pretrain = True
        logger.info(f"Training from latest: {latest_model}")
    else:
        try:
            minio_client.fget_object(MINIO_MODEL_TRAINED, request.pretrain, MODEL_CHECKPOINT)
        except Exception as exc:
            is_training = False
            raise HTTPException(status_code=500, detail=str(exc))
        use_pretrain = True
        logger.info(f"Training from : {request.pretrain}")

    objects = minio_client.list_objects(MINIO_DATA_LABELED, recursive=True)

    dataframes = []
    
    for obj in objects:
        logger.info(f'Data found: {obj.object_name}')
        try:
            response = minio_client.get_object(MINIO_DATA_LABELED, obj.object_name)
            file_data = BytesIO(response.read())
            response.close()
            response.release_conn()
            df = pd.read_csv(file_data)
            df = df[['text', 'label']]
            dataframes.append(df)
        except Exception as exc:
            is_training = False
            raise HTTPException(status_code=500, detail=str(exc))
    if len(dataframes) == 0:
        is_training = False
        raise HTTPException(status_code=500, detail="Can not get any data file")
    df = pd.concat(dataframes)
    df = df.dropna()

    logger.info("Pull data done")

    # Tien xu ly: Phan doan tu 
    try:
        df['text'] = df.apply(preprocess_row, axis=1)
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        is_training = False
        raise HTTPException(status_code=500, detail=str(e))
    df = df.dropna()
    label_counts = df['label'].value_counts()
    logger.info(f'Data label count: {label_counts}')
    
    # Sử dụng train_test_split để phân chia dữ liệu
    train_df, test_df = train_test_split(df, test_size=TEST_RATIO, stratify=df['label'], random_state=42)
    train_df = train_df.reset_index(drop=True)
    test_df = test_df.reset_index(drop=True)

    # Kiểm tra kết quả
    logger.info(f"Train DataFrame size: {train_df.shape}")
    logger.info(f"Test DataFrame size: {test_df.shape}")

    # We will use Kfold later
    skf = StratifiedKFold(n_splits=K_FOLD)
    for fold, (_, val_) in enumerate(skf.split(X=train_df, y=train_df.label)):
        train_df.loc[val_, "kfold"] = fold

    tokenizer = AutoTokenizer.from_pretrained(PHOBERTBASE_DIR, local_files_only=True, use_fast=False)
    # This ensures that any random number generation in NumPy, PyTorch (CPU and GPU), and cuDNN will be consistent
    seed_everything(86)

    stop_training_flag.clear()
    train_thread = threading.Thread(target=train_model, args=(skf, train_df, tokenizer, use_pretrain))
    train_thread.start()

    return TrainingResponse(status="Training started", message="Tracking and visualizing metrics on TensorBoard UI: http://localhost:6006/")

@app.post("/stop-training", response_model=TrainingResponse)
async def stop_training():
    global stop_training_flag, train_thread, is_training

    if not train_thread or not train_thread.is_alive():
        raise HTTPException(status_code=400, detail="No training in progress")

    stop_training_flag.set()
    train_thread.join()  # Wait for the thread to finish
    release_training_vram()
    is_training = False
    return TrainingResponse(status="Training stopped", message="VRam GPU are released")