# DNNGP3 tuning hyperparameters script

import os

import re

import json

import subprocess

import numpy as np

import nevergrad as ng

import tensorflow as tf

# The script needs to set parameters in three places, one is the #10 directory location, the second is the #21 hyperparameter search space, and the third is the #49 DNNGP native command.

# Set priorities in descending order, except for the directory, the default parameters are sufficient for most requests.

# Define directories and file paths

output_dir = r'..\Output_files'

pkl_file = r"..\Input_files\wheat599_pc95.pkl"

budget = 200  # Optimize the number of script iterations

alpha = 0.7  # Adjust this weight to balance the importance of mean and variance in the optimization process

beta = 0.1  # This parameter is adjusted to control the nonlinear effect of the deviation in the optimization process

cvs = 10 # K-fold cross-validation

pkl_dir = os.path.dirname(pkl_file)

# Obtain all tsv files in the directory where the pkl file resides

tsv_files = [f for f in os.listdir(pkl_dir) if f.endswith('.tsv')]

def check_gpu_available():

    """Check and display GPU availability information"""

    gpus = tf.config.list_physical_devices('GPU')

    if gpus:

        print("🎉 GPU is available!")

        for idx, gpu in enumerate(gpus):

            print(f"[Device {idx}]")

            print(f"  Name: {gpu.name}")

            try:

                details = tf.config.experimental.get_device_details(gpu)

                print(f"  Compute Capability: {details.get('compute_capability')}")

                print(f"  Device Type: {details.get('device_type', 'N/A')}")

            except AttributeError:

                print("  Unable to retrieve detailed device information (may require TensorFlow version upgrade)")

        return True

    else:

        print("⚠️ No GPU detected, will use CPU")

        return False

check_gpu_available()

# Define hyperparameters search space (see https://github.com/facebookresearch/nevergrad)

instr = ng.p.Instrumentation(

    batch_size=ng.p.Scalar(lower=32, upper=1024).set_integer_casting(),

    lr=ng.p.Log(lower=1e-4, upper=1),

    patience=ng.p.Scalar(lower=10, upper=50).set_integer_casting(),

    dropout1=ng.p.Log(lower=0.01, upper=0.9),

    dropout2=ng.p.Log(lower=0.01, upper=0.9),

    earlystopping=ng.p.Scalar(lower=50, upper=100).set_integer_casting()

)

# Define a function to extract statistic values

def extract_statistics(output):

    statistics = re.findall(r'statistic=([-+]?[0-9]*\.?[0-9]+)', output)

    if not statistics:

        return 0.0

    statistic_values = float(statistics[0])

    return statistic_values

# Define the objective function

def objective(batch_size: int, lr: float, patience: int, dropout1: float, dropout2: float, earlystopping: int, tsv_file: str):

    accuracies = []

    print('batch:',batch_size, 'lr:', lr, 'patience:', patience, 'dropout1:', dropout1, 'dropout2:', dropout2, 'earlystopping:', earlystopping, 'tsv_file:', tsv_file)

    for part in range(1, cvs + 1):

        command = f"python ../Scripts/dnngp_runner.py --batch_size {batch_size} --epoch 10000 --lr {lr} --patience {patience} --dropout1 {dropout1} --dropout2 {dropout2} --earlystopping {earlystopping} --cv {cvs} --part {part} --snp {pkl_file} --pheno {os.path.join(pkl_dir, tsv_file)} --output {output_dir}"

        print(command)

        p = subprocess.Popen(command, shell=True,

                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        output, error = p.communicate()

        # Decode output

        output_str = output.decode(errors='ignore')

        error_str = error.decode(errors='ignore')

        if error_str:

            print("Error Output:", error_str)

        accuracy = extract_statistics(output_str)

        accuracies.append(accuracy)

    print("Statistic values for all folds", accuracies)

    mean_accuracy = np.mean(accuracies) if accuracies else 0.0

    var_accuracy = np.var(accuracies) if accuracies else 0.0

# Use a weighted combination to balance the mean and variance

# Use a nonlinear transform to adjust for the effect of variance

    combined_metric = alpha * mean_accuracy - \

        (1 - alpha) * np.exp(beta * var_accuracy)

    return -combined_metric

# Record the best parameters and results for each tsv file

best_params_per_tsv = {}

for tsv_file in tsv_files:

    print(f"Optimizing for TSV file: {tsv_file}")

    # Use Nevergrad's optimizer

    optimizer = ng.optimizers.NGOpt(parametrization=instr, budget=budget)

    # Execution optimization procedure

    recommendation = optimizer.minimize(

        lambda *args, **kwargs: objective(*args, **kwargs, tsv_file=tsv_file)

    )

    # Output optimum parameter

    print(f"Best parameters for {tsv_file}:", recommendation.value)

    best_params_per_tsv[tsv_file] = recommendation.value

# Output best_params_per_tsv to a JSON file

output_json_file = os.path.join(pkl_dir, 'best_params_per_tsv.json')

with open(output_json_file, 'w') as file:

    json.dump(best_params_per_tsv, file, indent=4)

print(f"Best parameters saved to {output_json_file}")