import numpy as np
import os
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from generator import AugmentedImageSequence
import torch.optim as optim
from tensorflow.keras.models import model_from_json # type: ignore

def get_dataloader(csv, batch_size, FLAGS, tokenizer_wrapper, augmenter=None):
    """
    Replaces the TensorFlow enqueuer with PyTorch DataLoader.
    """
    data_generator = AugmentedImageSequence(
        dataset_csv_file=csv,
        class_names=FLAGS.csv_label_columns,
        tokenizer_wrapper=tokenizer_wrapper,
        source_image_dir=FLAGS.image_directory,
        batch_size=batch_size,
        target_size=FLAGS.image_target_size,
        augmenter=augmenter,
        shuffle_on_epoch_end=True,
    )
    dataloader = DataLoader(data_generator, shuffle=True, num_workers=0)
    return dataloader, data_generator.steps


def get_layers(layer_sizes, activation='relu'):
    """
    Builds a list of layers in PyTorch based on specified sizes.
    Dropout layers are specified with values < 1, Dense (Linear) layers otherwise.
    """
    layers = []
    activation_fn = getattr(nn, activation.capitalize(), nn.ReLU)  # Set default activation to ReLU if none specified
    for layer_size in layer_sizes:
        if layer_size < 1:
            layers.append(nn.Dropout(layer_size))
        else:
            layers.append(nn.Linear(in_features=layer_size, out_features=layer_size))
            layers.append(activation_fn())
    return nn.Sequential(*layers)  # Return as a sequential module for easy stacking



def get_optimizer(optimizer_type, learning_rate, lr_decay=0):
    optimizer_class = getattr(optim, optimizer_type)
    dummy_param = torch.nn.Parameter(torch.empty(0)) 
    optimizer = optimizer_class(params = [dummy_param], lr=learning_rate, weight_decay=lr_decay)
    return optimizer


def load_model(load_path, model_name):
    path = os.path.join(load_path, model_name)

    # load json and create model
    json_file = open('{}.json'.format(path), 'r')
    loaded_model_json = json_file.read()
    json_file.close()
    loaded_model = model_from_json(loaded_model_json)
    # # load weights into new model
    loaded_model.load_weights("{}.h5".format(path))
    print("Loaded model from disk")
    return loaded_model