import argparse
import json
import os
import traceback

import pandas as pd
from ab.nn.util.Util import release_memory, uuid4, read_py_file_as_string

from ab.gpt.util.Const import epoch_dir, synth_dir, new_nn_file, nngpt_dir
from ab.gpt.util.NNEval import NNEval
from ab.gpt.util.Util import verify_nn_code, copy_to_lemur

# --- Default Evaluation Parameters ---
# TODO: CARRY ALL THIS STUFF TO CONST.PY
# These will be used as defaults for argparse arguments
DEFAULT_NN_TRAINING_EPOCHS = 1  # How many epochs to train the altered NN for evaluation
DEFAULT_TASK = 'img-classification'
DEFAULT_DATASET = 'cifar-10'
DEFAULT_METRIC = 'acc'

# Default hyperparameters. 'epoch' will be overridden.
# TODO: DETECT EPOCH AUTOMATICALLY
DEFAULT_LR = 0.01
DEFAULT_BATCH_SIZE = 64
DEFAULT_DROPOUT = 0.2
DEFAULT_MOMENTUM = 0.9
DEFAULT_TRANSFORM = 'norm_256_flip'  # A common default, used by NNEval if prm is None
DEFAULT_NN_NAME_PREFIX = None

def evaluate_altered_models(args: argparse.Namespace):
    """
    Evaluates neural networks generated by NNAlter.py.

    :param args: Parsed command-line arguments.
    """
    print(f"Starting evaluation of altered NNs...")
    print(f"NNAlter run epochs to scan: {args.nn_alter_epochs}")
    print(f"Each altered NN will be trained for: {args.nn_train_epochs} epochs for evaluation.")
    print(f"Base task: {args.task}, Base dataset: {args.dataset}, Base metric: {args.metric}")
    print(f"Base Hyperparameters for NNEval (before df override):")
    print(f"  LR: {args.lr}, Batch Size: {args.batch_size}, Dropout: {args.dropout}, Momentum: {args.momentum}, Transform: {args.transform}")
    print(f"Save to DB: {args.save_to_db}")

    base_nngpt_path = nngpt_dir  # out/nngpt
    if not args.nn_alter_epochs:
        args.nn_alter_epochs = len(os.listdir(epoch_dir()))

    for i in range(args.nn_alter_epochs):
        # Path to the output of one NNAlter.py epoch (e.g., out/nngpt/llm/epoch/A0)
        current_alter_epoch_path = epoch_dir(i)  # This already uses nngpt_dir as base

        # Path to the synthesized models for that NNAlter epoch (e.g., .../A0/synth_nn)
        models_base_dir = synth_dir(current_alter_epoch_path)

        if not models_base_dir.exists():
            print(f"Directory {models_base_dir} for NNAlter epoch {i} not found. Skipping.")
            continue

        print(f"\n--- Scanning NNAlter Epoch Directory: {current_alter_epoch_path} ---")
        print(f"--- Synthesized Models Directory: {models_base_dir} ---")

        for model_folder_name in os.listdir(models_base_dir):
            model_dir_path = models_base_dir / model_folder_name
            if not model_dir_path.is_dir():
                continue

            code_file_path = model_dir_path / new_nn_file
            df_file_path = model_dir_path / 'dataframe.df'  # Original model's metadata

            if not code_file_path.exists():
                print(f"Code file {new_nn_file} not found in {model_dir_path}. Skipping.")
                continue

            print(f"\n--- Evaluating Model: {model_dir_path.relative_to(base_nngpt_path)} ---")

            if not verify_nn_code(model_dir_path, code_file_path):
                print(f"Code verification failed for {code_file_path}. Skipping evaluation.")
                with open(model_dir_path / 'eval_verification_failed.txt', 'w') as f:
                    f.write("Initial code verification failed.")
                continue

            # Initialize task, dataset, metric, and prm from command-line arguments (or their defaults)
            task = args.task
            dataset = args.dataset
            metric = args.metric
            # This prm structure is consistent with LEMUR dataset's expectations for model training
            prm = {
                'lr': args.lr,
                'batch': args.batch_size,
                'dropout': args.dropout,
                'momentum': args.momentum,
                'transform': args.transform,  # Default transform from CLI
                # 'epoch' will be set explicitly later
            }
            prefix_for_db = "AlteredNN"  # Default prefix
            origdf = None
            if df_file_path.exists():
                try:
                    origdf = pd.read_pickle(df_file_path)
                    # Override with values from dataframe.df if they exist
                    task = origdf.get('task', task)
                    dataset = origdf.get('dataset', dataset)
                    metric = origdf.get('metric', metric)

                    original_prm_from_df = origdf.get('prm')
                    if isinstance(original_prm_from_df, dict):
                        # Update prm with values from df, df values take precedence.
                        # This will update lr, batch, dropout, momentum, transform if they exist in original_prm_from_df,
                        # and also add any other model-specific hyperparameters from the original model.
                        prm.update(original_prm_from_df)

                    prefix_for_db = origdf.get('nn', prefix_for_db).split('-')[0]
                    print(f"  Loaded metadata from dataframe.df: task={task}, dataset={dataset}, metric={metric}")
                except Exception as e:
                    print(f"  Error loading dataframe.df from {df_file_path}: {e}. Using command-line/default parameters for task, dataset, metric, and prm structure.")
            else:
                print(f"  No dataframe.df found. Using command-line/default evaluation parameters.")

            # Crucial: set training epochs for this evaluation from args.nn_train_epochs
            # This overrides any 'epoch' value that might have come from original_prm_from_df
            prm['epoch'] = args.nn_train_epochs
            print(f"  Final parameters for NNEval: {prm}")
            print(f"  Task: {task}, Dataset: {dataset}, Metric: {metric}, Prefix: {prefix_for_db}")

            try:
                evaluator = NNEval(
                    model_source_package=str(model_dir_path),
                    task=task,
                    dataset=dataset,
                    metric=metric,
                    prm=prm,  # Pass the constructed prm
                    save_to_db=args.save_to_db,
                    prefix=prefix_for_db,
                    save_path=model_dir_path
                )
                eval_results = evaluator.evaluate(code_file_path)

                print(f"  Evaluation results for {model_folder_name}: {eval_results}")

                eval_info_data = {
                    "eval_args": evaluator.get_args(),  # This will show the prm used by NNEval
                    "eval_results": eval_results,
                    "cli_args": vars(args)
                }
                with open(model_dir_path / 'eval_info_altered.json', 'w+') as f:
                    json.dump(eval_info_data, f, indent=4, default=str)

                nn_name = uuid4(read_py_file_as_string(code_file_path))

                if args.nn_name_prefix:
                    nn_name = args.nn_name_prefix + '-' + nn_name
                copy_to_lemur(origdf, model_dir_path, nn_name)

            except Exception as e:
                error_msg = f"Error evaluating model {model_folder_name}: {e}"
                print(f"  {error_msg}")
                detailed_error = traceback.format_exc()
                print(detailed_error)
                with open(model_dir_path / 'eval_error_altered.txt', 'w+') as f:
                    f.write(f"{error_msg}\n\n{detailed_error}")
            finally:
                release_memory()


def main():
    parser = argparse.ArgumentParser(description="Evaluate Neural Networks generated by NNAlter.py.")
    parser.add_argument(
        '-nae', '--nn_alter_epochs', type=int, default=None,
        help="Number of epochs NNAlter.py was run for (e.g., if NNAlter's -e was 8, use 8 here)."
    )
    parser.add_argument(
        '-nte', '--nn_train_epochs', type=int, default=DEFAULT_NN_TRAINING_EPOCHS,
        help=f"Number of epochs to train each altered NN during evaluation (default: {DEFAULT_NN_TRAINING_EPOCHS})."
    )

    # Configurable evaluation parameters
    parser.add_argument(
        '--task', type=str, default=DEFAULT_TASK,
        help=f"Default task for NNEval if not in dataframe.df (default: {DEFAULT_TASK})."
    )
    parser.add_argument(
        '--dataset', type=str, default=DEFAULT_DATASET,
        help=f"Default dataset for NNEval if not in dataframe.df (default: {DEFAULT_DATASET})."
    )
    parser.add_argument(
        '--metric', type=str, default=DEFAULT_METRIC,
        help=f"Default metric for NNEval if not in dataframe.df (default: {DEFAULT_METRIC})."
    )

    # Configurable hyperparameters (part of prm dictionary for NNEval)
    parser.add_argument(
        '--lr', type=float, default=DEFAULT_LR,
        help=f"Learning rate for NNEval if not in dataframe.df's prm (default: {DEFAULT_LR})."
    )
    parser.add_argument(
        '--batch_size', type=int, default=DEFAULT_BATCH_SIZE,
        help=f"Batch size for NNEval if not in dataframe.df's prm (default: {DEFAULT_BATCH_SIZE}). Stored as 'batch' in prm."
    )
    parser.add_argument(
        '--dropout', type=float, default=DEFAULT_DROPOUT,
        help=f"Dropout rate for NNEval if not in dataframe.df's prm (default: {DEFAULT_DROPOUT})."
    )
    parser.add_argument(
        '--momentum', type=float, default=DEFAULT_MOMENTUM,
        help=f"Momentum for NNEval if not in dataframe.df's prm (default: {DEFAULT_MOMENTUM})."
    )
    parser.add_argument(
        '--transform', type=str, default=DEFAULT_TRANSFORM,
        help=f"Default transform for NNEval if not in dataframe.df's prm (default: {DEFAULT_TRANSFORM}). Stored as 'transform' in prm."
    )

    # Other NNEval options
    parser.add_argument(
        '--save_to_db', action=argparse.BooleanOptionalAction, default=True,
        help="Whether to save evaluation results to the database (enables with --save-to-db, disables with --no-save-to-db; default: enabled)."
    )

    parser.add_argument(
        '--nn_name_prefix', type=str, default=DEFAULT_NN_NAME_PREFIX,
        help=f"Default neural network name prefix (default: {DEFAULT_NN_NAME_PREFIX})."
    )

    args = parser.parse_args()

    evaluate_altered_models(args)

if __name__ == "__main__":
    main()

