import matplotlib.pyplot as plt
import numpy as np
from scipy import stats
from sklearn.model_selection import train_test_split
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint  # type: ignore
from model_unet import unet
from data_tools import scaled_in, scaled_ou

def training(path_save_spectrogram, weights_path, name_model, training_from_scratch, epochs, batch_size):
    """ This function will read noisy voice and clean voice spectrograms created by data_creation mode,
    and train a Unet model on this dataset for epochs and batch_size specified. It saves best models to disk regularly
    If training_from_scratch is set to True it will train from scratch, if set to False, it will train
    from weights (name_model) provided in weights_path
    """
    # Load noisy voice & clean voice spectrograms created by data_creation mode
    X_in = np.load(path_save_spectrogram +'noisy_voice_amp_db'+".npy")
    X_ou = np.load(path_save_spectrogram +'voice_amp_db'+".npy")
    # Model of noise to predict
    X_ou = X_in - X_ou

    # Check distribution
    print(stats.describe(X_in.reshape(-1,1)))
    print(stats.describe(X_ou.reshape(-1,1)))

    # To scale between -1 and 1
    X_in = scaled_in(X_in)
    X_ou = scaled_ou(X_ou)

    # Check shape of spectrograms
    print(X_in.shape)
    print(X_ou.shape)
    # Check new distribution
    print(stats.describe(X_in.reshape(-1,1)))
    print(stats.describe(X_ou.reshape(-1,1)))


    # Reshape for training
    X_in = X_in[:,:,:]
    X_in = X_in.reshape(X_in.shape[0], X_in.shape[1], X_in.shape[2], 1)
    X_ou = X_ou[:,:,:]
    X_ou = X_ou.reshape(X_ou.shape[0], X_ou.shape[1], X_ou.shape[2], 1)

    X_train, X_test, y_train, y_test = train_test_split(X_in, X_ou, test_size=0.10, random_state=42)

    # If training from scratch
    if training_from_scratch:
        generator_nn = unet()
    # If training from pre-trained weights
    else:
        generator_nn = unet(pretrained_weights=weights_path + name_model + '.h5')

    # Save best models to disk during training
    checkpoint = ModelCheckpoint(weights_path + '/model_best.keras', verbose=1, monitor='val_loss', save_best_only=True, mode='auto')

    # Training
    history = generator_nn.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, shuffle=True, 
                               callbacks=[checkpoint], verbose=1, validation_data=(X_test, y_test))

    # Save the full model after training
    generator_nn.save(weights_path + '/final_model.keras')  # Save the complete model

    # Plot training and validation loss (log scale)
    loss = history.history['loss']
    val_loss = history.history['val_loss']
    epochs = range(1, len(loss) + 1)

    plt.plot(epochs, loss, label='Training loss')
    plt.plot(epochs, val_loss, label='Validation loss')
    plt.yscale('log')
    plt.title('Training and validation loss')
    plt.legend()
    plt.show()

    print("Training complete. The model has been saved.")

