import csv
import numpy as np
import os
import pandas as pd
from torch import autocast
from tqdm import tqdm
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
from torch.utils.data import DataLoader
import argparse
from sklearn import metrics
from contextlib import nullcontext
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report

from datasets.emotion_dataset import get_training_set, get_test_set, get_full_training_set, get_ft_weighted_sampler
from models.mn.model import get_model as get_mobilenet
from models.dymn.model import get_model as get_dymn
from models.preprocess import AugmentMelSTFT
from helpers.init import worker_init_fn
from helpers.utils import NAME_TO_WIDTH, exp_warmup_linear_down


# EarlyStopping class
class EarlyStopping:
    def __init__(self, patience=5, delta=0, mode='loss'):
        """
        Early stopping to stop training when a monitored metric has stopped improving.

        Args:
            patience (int): How long to wait after last time the monitored metric improved.
            delta (float): Minimum change in the monitored metric to qualify as an improvement.
            mode (str): Metric to monitor, either 'loss' or 'acc'.
        """
        self.patience = patience
        self.delta = delta
        self.mode = mode
        self.best_score = None
        self.early_stop = False
        self.counter = 0

    def __call__(self, val_metric, model, path):
        """
        Args:
            val_metric (float): The current value of the monitored metric (val_loss or val_acc).
            model (torch.nn.Module): The model to save if the monitored metric improves.
            path (str): Path to save the model checkpoint.
        """
        if self.mode == 'acc':
            score = val_metric  # Higher val_acc is better
        else:
            score = -val_metric  # Lower val_loss is better

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(model, path)
        elif score < self.best_score + self.delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(model, path)
            self.counter = 0

    def save_checkpoint(self, model, path):
        """Save the model when the monitored metric improves."""
        #Save only weights
        #torch.save(model.state_dict(), path)
        
        #Save structure and weights
        torch.save(model, path)

def _mel_forward(x, mel):
    old_shape = x.size()
    x = x.reshape(-1, old_shape[2])
    x = mel(x)
    x = x.reshape(old_shape[0], old_shape[1], x.shape[1], x.shape[2])
    return x



def evaluate_metrics_with_cm(model, mel, dataloader, device, loss_function):
    """
    Advance version of evaluate_metrics
    """
    model.eval()
    mel.eval()
    eval_loss = 0
    correct, total = 0, 0
    all_predictions = []
    all_labels = []
    
    pbar = tqdm(dataloader)
    pbar.set_description("Final Evaluation")

    with torch.no_grad():
        for batch in pbar:
            x, _, y = batch
            x, y = x.to(device), y.to(device)

            # Convert labels to single-label format (if needed)
            if len(y.shape) > 1:  # check if y is one-hot encoding
                y = torch.argmax(y, dim=1)  # Convert to index

            x = _mel_forward(x, mel)
            y_hat, _ = model(x)
            loss = loss_function(y_hat, y)
            eval_loss += loss.item()

            predictions = torch.argmax(y_hat, dim=1)
            correct += (predictions == y).sum().item()
            total += y.size(0)
            
            all_predictions.extend(predictions.cpu().numpy())
            all_labels.extend(y.cpu().numpy())

    # Compute confusion matrix
    cm = confusion_matrix(all_labels, all_predictions)
    accuracy = correct / total
    
    return eval_loss / len(dataloader), accuracy, cm, all_labels, all_predictions

def save_confusion_matrix(cm, class_names, save_dir, model_name, timestamp, accuracy):
    """
    Draw and save confusion matrix
    """
    plt.figure(figsize=(max(12, len(class_names) * 0.4), max(10, len(class_names) * 0.4)))
    
    # Create Heatmap
    sns.heatmap(cm, 
                annot=True if len(class_names) <= 20 else False,  # Visualize number if number of class is small
                fmt='d', 
                cmap='Blues',
                xticklabels=class_names,
                yticklabels=class_names,
                cbar_kws={'label': 'Count'})
    
    plt.title(f'Confusion Matrix - {model_name}\nAccuracy: {accuracy:.4f}', 
              fontsize=16, fontweight='bold')
    plt.xlabel('Predicted Label', fontsize=12)
    plt.ylabel('True Label', fontsize=12)
    
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.tight_layout()
    
    # Draw
    cm_image_path = os.path.join(save_dir, f'confusion_matrix_{model_name}_{timestamp}.png')
    plt.savefig(cm_image_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    # Save confusion matrix in  CSV
    cm_csv_path = os.path.join(save_dir, f'confusion_matrix_{model_name}_{timestamp}.csv')
    cm_df = pd.DataFrame(cm, index=class_names, columns=class_names)
    cm_df.to_csv(cm_csv_path)
    
    print(f"Confusion matrix saved as image: {cm_image_path}")
    print(f"Confusion matrix saved as CSV: {cm_csv_path}")
    
    return cm_image_path, cm_csv_path
def generate_classification_report(y_true, y_pred, class_names, save_dir, model_name, timestamp):
    """
    Create and  save detail classification report
    """
    # Create classification report
    report = classification_report(y_true, y_pred, 
                                 target_names=class_names, 
                                 output_dict=True)
    
    # convert to DataFrane
    report_df = pd.DataFrame(report).transpose()
    
    # Save into csv
    report_path = os.path.join(save_dir, f'classification_report_{model_name}_{timestamp}.csv')
    report_df.to_csv(report_path)
    
    # print report summary
    print("\n=== CLASSIFICATION REPORT ===")
    print(f"Overall Accuracy: {report['accuracy']:.4f}")
    print(f"Macro Average Precision: {report['macro avg']['precision']:.4f}")
    print(f"Macro Average Recall: {report['macro avg']['recall']:.4f}")
    print(f"Macro Average F1-score: {report['macro avg']['f1-score']:.4f}")
    print(f"Weighted Average F1-score: {report['weighted avg']['f1-score']:.4f}")
    
    print(f"Classification report saved: {report_path}")
    return report_path

def evaluate_metrics(model, mel, dataloader, device, loss_function):
    model.eval()
    mel.eval()
    eval_loss = 0
    correct, total = 0, 0
    all_predictions = []
    all_labels = []
    
    pbar = tqdm(dataloader)
    pbar.set_description("Validating")

    with torch.no_grad():
        for x, _, y in pbar:
            x, y = x.to(device), y.to(device)

            # Convert labels to single-label format (if needed)
            if len(y.shape) > 1:  # Kiểm tra nếu y là one-hot encoding
                y = torch.argmax(y, dim=1)  # Chuyển sang dạng chỉ số

            x = _mel_forward(x, mel)
            y_hat, _ = model(x)
            loss = loss_function(y_hat, y)
            eval_loss += loss.item()

            predictions = torch.argmax(y_hat, dim=1)
            correct += (predictions == y).sum().item()
            total += y.size(0)
            
            # Lưu dự đoán và nhãn thực tế
            all_predictions.extend(predictions.cpu().numpy())
            all_labels.extend(y.cpu().numpy())

    # Tính toán ma trận nhầm lẫn
    cm = confusion_matrix(all_labels, all_predictions)
    
    # In ma trận nhầm lẫn
    print("Confusion Matrix:")
    print(cm)
    
    return eval_loss / len(dataloader), correct / total


# def evaluate_weighted_accuracy_with_confusion_matrix(model, mel, dataloader, device, loss_function, num_classes):
#     model.eval()
#     mel.eval()
#     eval_loss = 0
#     class_counts = np.zeros(num_classes, dtype=np.int32)
#     correct_counts = np.zeros(num_classes, dtype=np.int32)
#     total_samples = 0
#     y_true = []
#     y_pred = []

#     pbar = tqdm(dataloader)
#     pbar.set_description("Validating")

#     with torch.no_grad():
#         for x, _, y in pbar:
#             x, y = x.to(device), y.to(device)

#             # Convert labels to single-label format (if needed)
#             if len(y.shape) > 1:  # Nếu y là one-hot encoding
#                 y = torch.argmax(y, dim=1)  # Chuyển sang dạng chỉ số

#             x = _mel_forward(x, mel)
#             y_hat, _ = model(x)
#             loss = loss_function(y_hat, y)
#             eval_loss += loss.item()

#             predictions = torch.argmax(y_hat, dim=1)
#             total_samples += y.size(0)

#             # Cập nhật số lượng mẫu và số lượng đúng theo từng lớp
#             for i in range(num_classes):
#                 class_counts[i] += (y == i).sum().item()
#                 correct_counts[i] += ((predictions == i) & (y == i)).sum().item()

#             # Lưu trữ kết quả thực và dự đoán
#             y_true.extend(y.cpu().numpy())
#             y_pred.extend(predictions.cpu().numpy())

#     # Tính trọng số dựa trên số lượng mẫu
#     class_weights = np.where(class_counts > 0, 1.0 / class_counts, 0)
    
#     # Tính toán độ chính xác có trọng số
#     weighted_accuracy = np.sum(correct_counts * class_weights) / np.sum(class_weights)

#     # Tạo ma trận nhầm lẫn
#     conf_matrix = confusion_matrix(y_true, y_pred, labels=np.arange(num_classes))

#     # Trả về loss, weighted accuracy, confusion matrix và các thông tin bổ sung
#     return eval_loss / len(dataloader), weighted_accuracy, conf_matrix






def train(args):
    # Train Models from scratch or ImageNet pre-trained on AudioSet
    # PaSST ensemble (https://github.com/kkoutini/PaSST) stored in 'resources/passt_enemble_logits_mAP_495.npy'
    # can be used as a teacher.
    
    label_csv_path = args.label_csv_path
    if os.path.exists(label_csv_path) and label_csv_path.lower().endswith('.csv'):
        labels_df = pd.read_csv(label_csv_path)
        number_classes = len(labels_df)

        class_names = labels_df['display_name'].tolist()
    else:
        number_classes = 527
        class_names = [f'Class_{i}' for i in range(number_classes)]
        
    

    device = torch.device('cuda') if args.cuda and torch.cuda.is_available() else torch.device('cpu')

    #Create output_model 
    output_path = './output_model'
    output_best_checkpoint = './output_model/best_checkpoint'
    os.makedirs(output_path, exist_ok=True)
    os.makedirs(output_best_checkpoint, exist_ok=True)

    # Setup CSV logging
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_filename = f"./training_logs/training_metrics_{args.model_name}_{timestamp}.csv"
    results_dir = f"./training_results_{args.model_name}_{timestamp}"
    
    # Tạo thư mục nếu chưa có
    os.makedirs("./training_logs", exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)
    
    # Khởi tạo CSV file với headers
    csv_headers = [
        'epoch', 'train_loss', 'val_loss', 'val_accuracy', 
        'learning_rate', 'best_val_acc', 'time_per_epoch',
        'model_name', 'batch_size', 'optimizer_type'
    ]
    
    with open(csv_filename, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=csv_headers)
        writer.writeheader()


    # model to preprocess waveform into mel spectrograms
    mel = AugmentMelSTFT(n_mels=args.n_mels,
                         sr=args.resample_rate,
                         win_length=args.window_size,
                         hopsize=args.hop_size,
                         n_fft=args.n_fft,
                         freqm=args.freqm,
                         timem=args.timem,
                         fmin=args.fmin,
                         fmax=args.fmax,
                         fmin_aug_range=args.fmin_aug_range,
                         fmax_aug_range=args.fmax_aug_range
                         )
    mel.to(device)
    # load prediction model
    model_name = args.model_name
    pretrained_name = model_name if args.pretrained else None
    width = NAME_TO_WIDTH(model_name) if model_name and args.pretrained else args.model_width
    if model_name.startswith("dymn"):
        model = get_dymn(num_classes = number_classes, width_mult=width, pretrained_name=pretrained_name,
                         strides=args.strides, pretrain_final_temp=args.pretrain_final_temp)
    else:
        model = get_mobilenet(num_classes = number_classes, width_mult=width, pretrained_name=pretrained_name,
                              strides=args.strides, head_type=args.head_type, se_dims=args.se_dims)
    model.to(device)

    # dataloader
    dl = DataLoader(dataset=get_training_set(resample_rate=args.resample_rate, roll=args.roll, wavmix=args.wavmix,
                                                  gain_augment=args.gain_augment),
                    sampler=get_ft_weighted_sampler(args.epoch_len, num_classes = number_classes),  # sampler important to balance classes
                    worker_init_fn=worker_init_fn,
                    num_workers=args.num_workers,
                    batch_size=args.batch_size)

    # evaluation loader
    eval_dl = DataLoader(dataset=get_test_set(resample_rate=args.resample_rate),
                         worker_init_fn=worker_init_fn,
                         num_workers=args.num_workers,
                         batch_size=args.batch_size)

    
    optimizer_type = "Adam"
    if args.adamw:
        optimizer_type = "AdamW"
        # optimizer & scheduler
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.max_lr, weight_decay=args.weight_decay)
    else:
        # optimizer & scheduler
        optimizer = torch.optim.Adam(model.parameters(), lr=args.max_lr, weight_decay=args.weight_decay)

    #scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=3, verbose=True)

    # phases of lr schedule: exponential increase, constant lr, linear decrease, fine-tune
    schedule_lambda = \
        exp_warmup_linear_down(args.warm_up_len, args.ramp_down_len, args.ramp_down_start, args.last_lr_value)
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule_lambda)

    # Loss function for single-label classification
    loss_function = nn.CrossEntropyLoss(label_smoothing = args.label_smoothing)

    # Early stopping
    early_stopping = EarlyStopping(patience=args.patience, delta=args.delta, mode=args.early_stop_mode)

    best_val_acc = 0

    for epoch in range(args.n_epochs):
        epoch_start_time = datetime.now()
        mel.train()
        model.train()

        train_loss = []

        pbar = tqdm(dl)
        pbar.set_description(f"Epoch {epoch + 1}/{args.n_epochs}")

        # in case of DyMN: update DyConv temperature
        if hasattr(model, "update_params"):
            model.update_params(epoch)

        for batch in pbar:
            x, _, y, _ = batch

            # Convert labels to single-label format (if needed)
            if len(y.shape) > 1:  # Kiểm tra nếu y là one-hot encoding
                y = torch.argmax(y, dim=1)  # Chuyển sang dạng chỉ số

            x, y = x.to(device), y.to(device)
            x = _mel_forward(x, mel)

            y_hat, _ = model(x)
            loss = loss_function(y_hat, y)

            # Update Model
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss.append(loss.item())
            pbar.set_postfix(train_loss=sum(train_loss) / len(train_loss))

        # Calculate epoch metrics
        epoch_train_loss = sum(train_loss) / len(train_loss)
        current_lr = optimizer.param_groups[0]['lr']

        # Evaluate the model
        val_loss, val_acc = evaluate_metrics(model, mel, eval_dl, device, loss_function)

        # Calculate time per epoch
        epoch_end_time = datetime.now()
        time_per_epoch = (epoch_end_time - epoch_start_time).total_seconds()
        
        #val_loss, val_acc,  conf_matrix = evaluate_weighted_accuracy_with_confusion_matrix(model, mel, eval_dl, device, loss_function, number_classes)
        print(f"Epoch {epoch + 1}/{args.n_epochs}, Train Loss: {epoch_train_loss:.4f}, "
              f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_acc:.4f}, "
              f"LR: {current_lr:.6f}, Time: {time_per_epoch:.2f}s")
        #print(conf_matrix)

        # Save best checkpoint
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_model_path = f"{output_best_checkpoint}/{model_name}{str(width).replace('.', '')}_acc_{best_val_acc:.4f}.pt"
            ##Save only weights
            #torch.save(model.state_dict(), best_model_path)
            
            ##Save structure and weights
            torch.save(model, best_model_path)
        if args.save_per_epoch:
            name = f"{model_name}{str(width).replace('.', '')}_as_epoch_{epoch}_val_loss_{val_loss:.4f}_val_acc_{val_acc:.4f}.pt"
            # Save only weights
            # torch.save(model.state_dict(), os.path.join("./output_model", name))
            
            ###Save structure and weights
            torch.save(model, os.path.join(output_path, name))

        # Lưu metrics vào CSV
        epoch_metrics = {
            'epoch': epoch + 1,
            'train_loss': round(epoch_train_loss, 6),
            'val_loss': round(val_loss, 6),
            'val_accuracy': round(val_acc, 6),
            'learning_rate': round(current_lr, 8),
            'best_val_acc': round(best_val_acc, 6),
            'time_per_epoch': round(time_per_epoch, 2),
            'model_name': model_name,
            'batch_size': args.batch_size,
            'optimizer_type': optimizer_type
        }
        
        # Append to CSV
        with open(csv_filename, 'a', newline='', encoding='utf-8') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_headers)
            writer.writerow(epoch_metrics)

        # Update learning rate
        scheduler.step()

        # Early stopping check
        if early_stopping.mode == 'acc':
            early_stopping(val_acc, model, f'{output_path}/checkpoint.pt')
        else:
            early_stopping(val_loss, model, f'{output_path}/checkpoint.pt')

        if early_stopping.early_stop:
            print("Early stopping triggered.")
            # Ghi thông tin early stopping vào CSV
            with open(csv_filename, 'a', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow([f"Early stopping at epoch {epoch + 1}"])
            break

    # =================== POST-TRAINING ANALYSIS ===================
    print("\n" + "="*60)
    print("TRAINING COMPLETED - Generating Final Analysis")
    print("="*60)
    
    # Thực hiện evaluation cuối cùng để có confusion matrix
    print("Performing final evaluation...")
    final_val_loss, final_val_acc, cm, y_true, y_pred = evaluate_metrics_with_cm(
        model, mel, eval_dl, device, loss_function
    )
    
    print(f"Final Validation Loss: {final_val_loss:.4f}")
    print(f"Final Validation Accuracy: {final_val_acc:.4f}")
    
    # Vẽ và lưu confusion matrix
    cm_image_path, cm_csv_path = save_confusion_matrix(
        cm, class_names, results_dir, args.model_name, timestamp, final_val_acc
    )
    
    # Tạo classification report chi tiết
    report_path = generate_classification_report(
        y_true, y_pred, class_names, results_dir, args.model_name, timestamp
    )
    
    # Tạo summary file
    summary_path = os.path.join(results_dir, f'training_summary_{args.model_name}_{timestamp}.txt')
    with open(summary_path, 'w', encoding='utf-8') as f:
        f.write("TRAINING SUMMARY\n")
        f.write("="*50 + "\n\n")
        f.write(f"Model: {args.model_name}\n")
        f.write(f"Total Classes: {number_classes}\n")
        f.write(f"Training Duration: {timestamp}\n")
        f.write(f"Total Epochs: {epoch + 1}\n")
        f.write(f"Best Validation Accuracy: {best_val_acc:.4f}\n")
        f.write(f"Final Validation Accuracy: {final_val_acc:.4f}\n")
        f.write(f"Final Validation Loss: {final_val_loss:.4f}\n")
        f.write(f"Optimizer: {optimizer_type}\n")
        f.write(f"Batch Size: {args.batch_size}\n")
        f.write(f"Max Learning Rate: {args.max_lr}\n\n")
        f.write("FILES GENERATED:\n")
        f.write(f"- Training metrics: {csv_filename}\n")
        f.write(f"- Confusion matrix (image): {cm_image_path}\n")
        f.write(f"- Confusion matrix (CSV): {cm_csv_path}\n")
        f.write(f"- Classification report: {report_path}\n")
        f.write(f"- Training summary: {summary_path}\n")
    
    print(f"\nTraining completed successfully!")
    print(f"All results saved to: {results_dir}")
    print(f"Training metrics saved to: {csv_filename}")
    print(f"Summary saved to: {summary_path}")

def evaluate(args):
    model_name = args.model_name
    device = torch.device('cuda') if args.cuda and torch.cuda.is_available() else torch.device('cpu')


    print(f"Running AudioSet evaluation for model '{model_name}' on device '{device}'")
    if model_name.startswith("dymn"):
        model = get_dymn(width_mult=NAME_TO_WIDTH(model_name), pretrained_name=model_name,
                            strides=args.strides)
    else:
        model = get_mobilenet(width_mult=NAME_TO_WIDTH(model_name), pretrained_name=model_name,
                                strides=args.strides, head_type=args.head_type)
    model.to(device)
    model.eval()

    # model to preprocess waveform into mel spectrograms
    mel = AugmentMelSTFT(n_mels=args.n_mels,
                         sr=args.resample_rate,
                         win_length=args.window_size,
                         hopsize=args.hop_size,
                         n_fft=args.n_fft,
                         fmin=args.fmin,
                         fmax=args.fmax
                         )
    mel.to(device)
    mel.eval()

    dl = DataLoader(dataset=get_test_set(resample_rate=args.resample_rate),
                    worker_init_fn=worker_init_fn,
                    num_workers=args.num_workers,
                    batch_size=args.batch_size)

    targets = []
    predictions = []
    total_loss = []
    loss_function = nn.CrossEntropyLoss()
    for batch in tqdm(dl):
        x, _, y = batch
        # Convert labels to single-label format if needed
        if len(y.shape) > 1:  # If y is one-hot encoded
            y = torch.argmax(y, dim=1)
        
        x, y = x.to(device), y.to(device)
        # our models are trained in half precision mode (torch.float16)
        # run on cuda with torch.float16 to get the best performance
        # running on cpu with torch.float32 gives similar performance, using torch.bfloat16 is worse
        with autocast(device_type=device.type) if args.cuda else nullcontext():
            with torch.no_grad():
                x = _mel_forward(x, mel)
                y_hat, _ = model(x)
        
        # Compute loss
        loss = loss_function(y_hat, y)
        total_loss.append(loss.item())

        # Store predictions and targets
        preds = torch.argmax(y_hat, dim=1)
        predictions.extend(preds.cpu().numpy())
        targets.extend(y.cpu().numpy())

    # Calculate metrics
    avg_loss = np.mean(total_loss)
    accuracy = accuracy_score(targets, predictions)

    # Print evaluation results
    print(f"Results on test set for model: {model_name}")
    print(f"  Average Loss: {avg_loss:.4f}")
    print(f"  Accuracy: {accuracy:.4f}")

    # return avg_loss, accuracy

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Example of parser. ')

    # general
    parser.add_argument('--experiment_name', type=str, default="AudioSet")
    parser.add_argument('--train', action='store_true', default=False)
    parser.add_argument('--cuda', action='store_true', default=False)
    parser.add_argument('--batch_size', type=int, default=120)
    parser.add_argument('--num_workers', type=int, default=12)

    # evaluation
    parser.add_argument('--model_name', type=str, default="mn10_as")  # used also for training

    # training
    parser.add_argument('--pretrained', action='store_true', default=False)
    parser.add_argument('--pretrain_final_temp', type=float, default=30.0)  # for DyMN
    parser.add_argument('--model_width', type=float, default=1.0)
    parser.add_argument('--strides', nargs=4, default=[2, 2, 2, 2], type=int)
    parser.add_argument('--head_type', type=str, default="mlp")
    parser.add_argument('--se_dims', type=str, default="c")
    parser.add_argument('--n_epochs', type=int, default=200)
    parser.add_argument('--epoch_len', type=int, default=65536)
    parser.add_argument('--roll', action='store_true', default=False)
    parser.add_argument('--wavmix', action='store_true', default=False) #Cái này phù hợp với multi classification vì  trộn sóng âm thanh
    parser.add_argument('--gain_augment', type=int, default=0) #Augmentation (thay đổi mức độ âm thanh) VD Set = 5 thì tín hiệu thay đổi từ [-5,4]
    parser.add_argument('--label_smoothing', type=float, default=0.15)

    # optimizer
    parser.add_argument('--adamw', action='store_true', default=False)
    parser.add_argument('--weight_decay', type=float, default=0.01)
    # lr schedule
    parser.add_argument('--max_lr', type=float, default=0.0008)
    parser.add_argument('--warm_up_len', type=int, default=8)
    parser.add_argument('--ramp_down_start', type=int, default=80)
    parser.add_argument('--ramp_down_len', type=int, default=95)
    parser.add_argument('--last_lr_value', type=float, default=0.01)

    # early stopping
    parser.add_argument('--patience', type=int, default=50)
    parser.add_argument('--delta', type=int, default=0)
    parser.add_argument('--early_stop_mode', type=str, default='acc')

    # Save per Epoch
    parser.add_argument('--save_per_epoch', action='store_true', default=False)


    # preprocessing
    parser.add_argument('--resample_rate', type=int, default=32000)
    parser.add_argument('--window_size', type=int, default=800)
    parser.add_argument('--hop_size', type=int, default=320)
    parser.add_argument('--n_fft', type=int, default=1024)
    parser.add_argument('--n_mels', type=int, default=128)
    parser.add_argument('--freqm', type=int, default=0)
    parser.add_argument('--timem', type=int, default=0)
    parser.add_argument('--fmin', type=int, default=0)
    parser.add_argument('--fmax', type=int, default=None)
    parser.add_argument('--fmin_aug_range', type=int, default=10)
    parser.add_argument('--fmax_aug_range', type=int, default=2000)
    
    #Label_csv_path
    parser.add_argument('--label_csv_path', type=str, default='./metadata/class_labels_indices.csv')

    args = parser.parse_args()
    if args.train:
        train(args)
    else:
        evaluate(args)



def analyze_training_logs(csv_path):
    """
    Function to analyze training logs
    """
    df = pd.read_csv(csv_path)
    
    print("=== TRAINING SUMMARY ===")
    print(f"Total epochs: {len(df)}")
    print(f"Best validation accuracy: {df['val_accuracy'].max():.4f} at epoch {df.loc[df['val_accuracy'].idxmax(), 'epoch']}")
    print(f"Final training loss: {df['train_loss'].iloc[-1]:.4f}")
    print(f"Final validation loss: {df['val_loss'].iloc[-1]:.4f}")
    print(f"Total training time: {df['time_per_epoch'].sum():.2f} seconds")
    print(f"Average time per epoch: {df['time_per_epoch'].mean():.2f} seconds")
    
    # Draw
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    
    # Loss curves
    axes[0,0].plot(df['epoch'], df['train_loss'], label='Train Loss', color='blue')
    axes[0,0].plot(df['epoch'], df['val_loss'], label='Val Loss', color='red')
    axes[0,0].set_xlabel('Epoch')
    axes[0,0].set_ylabel('Loss')
    axes[0,0].set_title('Training and Validation Loss')
    axes[0,0].legend()
    axes[0,0].grid(True)
    
    # Accuracy curve
    axes[0,1].plot(df['epoch'], df['val_accuracy'], label='Val Accuracy', color='green')
    axes[0,1].set_xlabel('Epoch')
    axes[0,1].set_ylabel('Accuracy')
    axes[0,1].set_title('Validation Accuracy')
    axes[0,1].legend()
    axes[0,1].grid(True)
    
    # Learning rate
    axes[1,0].plot(df['epoch'], df['learning_rate'], label='Learning Rate', color='orange')
    axes[1,0].set_xlabel('Epoch')
    axes[1,0].set_ylabel('Learning Rate')
    axes[1,0].set_title('Learning Rate Schedule')
    axes[1,0].legend()
    axes[1,0].grid(True)
    
    # Time per epoch
    axes[1,1].bar(df['epoch'], df['time_per_epoch'], alpha=0.7, color='purple')
    axes[1,1].set_xlabel('Epoch')
    axes[1,1].set_ylabel('Time (seconds)')
    axes[1,1].set_title('Time per Epoch')
    axes[1,1].grid(True)
    
    plt.tight_layout()
    plt.show()
    
    return df


def load_and_visualize_confusion_matrix(cm_csv_path):
    """
    Load and visualize confusion matrix from csv
    """
    cm_df = pd.read_csv(cm_csv_path, index_col=0)
    
    plt.figure(figsize=(max(12, len(cm_df) * 0.4), max(10, len(cm_df) * 0.4)))
    
    sns.heatmap(cm_df.values, 
                annot=True if len(cm_df) <= 20 else False,
                fmt='d', 
                cmap='Blues',
                xticklabels=cm_df.columns,
                yticklabels=cm_df.index,
                cbar_kws={'label': 'Count'})
    
    plt.title('Confusion Matrix', fontsize=16, fontweight='bold')
    plt.xlabel('Predicted Label', fontsize=12)
    plt.ylabel('True Label', fontsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.show()
    
    return cm_df
