
"""
Tutorial bonus: Hyperparameter optimization for an hardware-aware SNN

In this example, we train a Neuronova hardware-aware spiking neural network for a binary
classification problem over audio data.

The model receives audio features over time and predicts one of two classes.
Each class is represented by a population of output neurons, so the final
prediction is obtained by population voting over output spikes.

The script demonstrates:

1. Building a multi-layer SNN deployable to hardware.
2. Initializing firing activity with frontend firing-rate initialization (and their hyper parameter tuning).
3. Applying fluctuation-based initialization (and their hyper parameter tuning).
4. Defining a composite loss with tuning of their hyper params:
   - population-based supervised loss
   - firing-rate regularization
5. Defining accuracy, precision, recall, and F1 metrics (that will be reported in the tensorboard).
6. Running hyperparameter optimization with Ray Tune through HPOWrapper.
"""

import os

import torch
import torch.nn as nn
from ray import tune

from nwavesdk.init.fluct_init import fluct_init
from nwavesdk.init.frontend_firing_init import frontend_firing_init
from nwavesdk.layers import H1v1Frontend, H1v1Layer, H1v1Synapse, prepare_net
from nwavesdk.loss import balanced_population_mse, firing_rate_target_mse_loss
from nwavesdk.metrics import accuracy_population
from nwavesdk.optim.hpo.HPOWrapper import HPOWrapper, train_one_trial
from nwavesdk.optim.hpo.loss_factory import CompositeLoss


def load_init_loader(loader_path):
    return torch.load(loader_path, weights_only=False)


class NET(nn.Module):
    """
        Since this is a binary classification task, we use:

        num_classes = 2

        Each class is represented by several output neurons. For example, with
        neurons_per_class = 32, the output layer has:

            2 * 32 = 64 output neurons

        The predicted class is obtained by counting spikes in each class population.

    """
    def __init__(
        self,
        in_ch,
        num_classes,
        neurons_per_class,
        dt,
        device,
        tau_l1,
        tau_l2,
        tau_l3,
        tau_l4,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.neurons_per_class = neurons_per_class
        self.out_neurons = num_classes * neurons_per_class
        self.dt = dt
        self.device_flag = device

        n1 = 13 # in this case using 13 filterbanks because the dataset was sampled at 16kHz, therefore last 3 channels are discarded because of Nyquist
        n2 = 64
        n3 = 64
        n4 = self.out_neurons         # Output neurons are split into class populations.

        self.s1 = H1v1Frontend( in_ch, n1, device=device)
        self.l1 = H1v1Layer( n1, taus=torch.full((n1,), tau_l1), dt=dt, device=device, layer_topology="FF")

        self.s2 = H1v1Synapse(n1, n2,device=device)
        self.l2 = H1v1Layer( n2, taus=torch.full((n2,), tau_l2), dt=dt, device=device, layer_topology="FF")

        self.s3 = H1v1Synapse( n2, n3, device=device)
        self.l3 = H1v1Layer( n3, taus=torch.full((n3,), tau_l3), dt=dt, device=device, layer_topology="FF")

        self.s4 = H1v1Synapse( n3, n4, device=device)
        self.l4 = H1v1Layer( n4, taus=torch.full((n4,), tau_l4), dt=dt, device=device, layer_topology="FF")

        self.frontend_stage = (self.s1, self.l1)
        self.layer_pairs = [
            (self.s2, self.l2),
            (self.s3, self.l3),
            (self.s4, self.l4),
        ]

        if device == "gpu":
            self.to_gpu()

    def to_gpu(self):
        for module in [self.s1, self.l1, self.s2, self.l2, self.s3, self.l3, self.s4, self.l4]:
            if hasattr(module, "to_gpu"):
                module.to_gpu()

    def forward_cpu(self, x):
        prepare_net(self)

        spk1, spk2, spk3, spk4 = [], [], [], []
        mem1, mem2, mem3, mem4 = [], [], [], []

        for t in range(x.shape[1]):
            q1 = self.s1(x[:, t, :])
            s1, m1 = self.l1(q1)

            q2 = self.s2(s1)
            s2, m2 = self.l2(q2)

            q3 = self.s3(s2)
            s3, m3 = self.l3(q3)

            q4 = self.s4(s3)
            s4, m4 = self.l4(q4)

            spk1.append(s1)
            mem1.append(m1)
            spk2.append(s2)
            mem2.append(m2)
            spk3.append(s3)
            mem3.append(m3)
            spk4.append(s4)
            mem4.append(m4)

        spikes_list = [
            torch.stack(spk1, dim=1),
            torch.stack(spk2, dim=1),
            torch.stack(spk3, dim=1),
            torch.stack(spk4, dim=1),
        ]
        membranes_list = [
            torch.stack(mem1, dim=1),
            torch.stack(mem2, dim=1),
            torch.stack(mem3, dim=1),
            torch.stack(mem4, dim=1),
        ]
        return spikes_list, membranes_list

    def forward(self, x):
        prepare_net(self)

        target_device = "cuda" if self.device_flag == "gpu" else "cpu"
        x = x.to(target_device, non_blocking=True).contiguous().to(torch.float32)

        q1 = self.s1(x)
        s1, m1 = self.l1(q1)

        q2 = self.s2(s1)
        s2, m2 = self.l2(q2)

        q3 = self.s3(s2)
        s3, m3 = self.l3(q3)

        q4 = self.s4(s3)
        s4, m4 = self.l4(q4)

        return [s1, s2, s3, s4], [m1, m2, m3, m4]
    


# ---------------------------------------------------------------------
# HPOWrapper setup
# ---------------------------------------------------------------------
# Here we register the builder functions defined above.
#
# HPOWrapper will call:
#
# - build_snn(...) to create the model
# - build_datasets(...) to load train/validation data
# - build_main_loss(...) to create the loss
# - build_*_metric(...) to create evaluation metrics
#
# This is the main extension point of the tutorial:
# users can keep the same HPO structure and replace only the builders
# with their own model, loss, dataset, or metric definitions.



# ---------------------------------------------------------------------
# Builder functions
# ---------------------------------------------------------------------
# HPOWrapper does not need to know the details of the model, dataset,
# loss, or metrics. Instead, it receives builder functions.
#
# Each builder follows a simple contract:
#
# - model_builder(model_cfg) -> nn.Module
# - dataset_builder(dataset_path, data_cfg) -> train_loader, val_loader
# - loss_builder(loss_cfg) -> loss object / callable
# - metric_builder(metric_cfg) -> metric function
#
# You are free to replace these functions with your own, as long as they
# follow the same kind of input/output contract expected by HPOWrapper.

def build_snn(model_cfg):
    """
        Build and initialize the SNN for one HPO trial.

        Contract:
            model_cfg -> initialized torch.nn.Module

        HPOWrapper calls this function at the beginning of each trial. The values
        inside model_cfg come from the parameters bound with hpo.bind_model_param().

        You can define your own model builder, for example to use a different
        architecture, as long as it receives a config dictionary and returns a
        ready-to-train PyTorch model.
    """
    model = NET(
        in_ch=model_cfg["in_ch"],
        num_classes=model_cfg["num_classes"],
        neurons_per_class=model_cfg["neurons_per_class"],
        dt=model_cfg["dt"],
        device=model_cfg.get("device", "gpu"),
        tau_l1=model_cfg["tau_l1"],
        tau_l2=model_cfg["tau_l2"],
        tau_l3=model_cfg["tau_l3"],
        tau_l4=model_cfg["tau_l4"],
    )

    init_loader = load_init_loader(model_cfg["init_train_loader_path"])

    if model_cfg.get("use_frontend_firing_init", True):
        frontend_firing_init(
            model,
            init_loader,
            target_fr=model_cfg.get("frontend_target_fr", 0.20),
            n_batches=model_cfg.get("frontend_n_batches", 2),
            epsilon=model_cfg.get("frontend_epsilon", 0.02),
            verbose=model_cfg.get("init_verbose", False),
        )

    if model_cfg.get("use_fluct_init", True):
        fluct_init(
            model,
            init_loader,
            xi_target=model_cfg["fluct_xi_target"],
            alpha=model_cfg["fluct_alpha"],
            n_batches=model_cfg["fluct_n_batches"],
            verbose=model_cfg.get("init_verbose", False),
        )

    return model


def balanced_pop_term(spk_out, y, config, **kwargs):
    """
        Supervised population-coding loss term.

        Contract:
            spk_out, y, config, **kwargs -> dict[str, torch.Tensor]

        This term compares the output spikes against the target label. Since each
        class is represented by a population of neurons, the loss encourages the
        correct class population to fire at correct_rate and the incorrect class
        population to fire at incorrect_rate.

        You can define your own loss term as long as it accepts the tensors passed
        by the trainer and returns a dictionary containing named loss values.
    """
    return {
        "mse": balanced_population_mse(
            spk_out,
            y,
            num_classes=config["num_classes"],
            neurons_per_class=config["neurons_per_class"],
            correct_rate=config["correct_rate"],
            incorrect_rate=config["incorrect_rate"],
        )
    }


def firing_rate_term(spikes_list, config, **kwargs):
    """
        Firing-rate regularization loss term.

        Contract:
            spikes_list, config, **kwargs -> dict[str, torch.Tensor]

        This term controls the activity level of the network. It is useful because
        very high firing rates can make training unstable.

        You can replace this with any other regularization term, provided it follows
        the same callable structure expected by CompositeLoss.
    """
    return {
        "firing_rate": firing_rate_target_mse_loss(
            spikes_list=spikes_list,
            offsets=config["offsets_fr"],
            multipliers=config["multipliers_fr"],
        )
    }

# ---------------------------------------------------------------------
# Population-voting utilities
# ---------------------------------------------------------------------
# These helper functions convert output spikes into class predictions.
#
# The output layer is organized as:
#
#     [class 0 neurons] [class 1 neurons]
#
# For each sample, we count spikes inside each class population and choose
# the class with the largest spike count.

def _popvote_pred(spk, num_classes):
    batch_size, timesteps, out_neurons = spk.shape
    assert out_neurons % num_classes == 0, "num_outputs must be a multiple of num_classes"
    neurons_per_class = out_neurons // num_classes
    counts = spk.view(batch_size, timesteps, num_classes, neurons_per_class).sum(dim=(1, 3))
    return counts.argmax(dim=1)


def precision_population(spk, targets, num_classes, positive_class=1):
    preds = _popvote_pred(spk, num_classes=num_classes).view(-1)
    targets = targets.view(-1)

    tp = ((preds == positive_class) & (targets == positive_class)).float().sum().item()
    fp = ((preds == positive_class) & (targets != positive_class)).float().sum().item()
    return tp / max(1.0, tp + fp)


def recall_population(spk, targets, num_classes, positive_class=1):
    preds = _popvote_pred(spk, num_classes=num_classes).view(-1)
    targets = targets.view(-1)

    tp = ((preds == positive_class) & (targets == positive_class)).float().sum().item()
    fn = ((preds != positive_class) & (targets == positive_class)).float().sum().item()
    return tp / max(1.0, tp + fn)


def f1_population(spk, targets, num_classes, positive_class=1):
    precision = precision_population(spk, targets, num_classes, positive_class)
    recall = recall_population(spk, targets, num_classes, positive_class)
    return 0.0 if (precision + recall) == 0.0 else 2.0 * precision * recall / (precision + recall)


def build_main_loss(loss_cfg):
    """
        Build the composite training loss.

        Contract:
            loss_cfg -> CompositeLoss

        HPOWrapper calls this function to create the loss used in each trial.
        The values inside loss_cfg come from hpo.bind_loss_param().

        This example combines:
        - a supervised population loss
        - a firing-rate regularization loss

        You can add, remove, or replace terms as long as the returned object behaves
        like a valid training loss.
    """
    return (
        CompositeLoss()
        .add_loss_term(
            "population",
            balanced_pop_term,
            config={
                "num_classes": 2,
                "neurons_per_class": 32,
                "correct_rate": loss_cfg["correct_rate"],
                "incorrect_rate": loss_cfg["incorrect_rate"],
            },
        )
        .add_loss_term(
            "regularization",
            firing_rate_term,
            config={
                "offsets_fr": loss_cfg["offsets_fr"],
                "multipliers_fr": loss_cfg["multipliers_fr"],
            },
        )
    )


def build_datasets(dataset_path, data_cfg):
    train_loader = torch.load(
        os.path.join(dataset_path, data_cfg.get("train_loader_file", "trainb32plus.pt")),
        weights_only=False,
    )
    val_loader = torch.load(
        os.path.join(dataset_path, data_cfg.get("val_loader_file", "valb32plus.pt")),
        weights_only=False,
    )
    return train_loader, val_loader


def build_acc_metric(metric_cfg):
    """
        The trainer will call the returned function during evaluation and log the result to TensorBoard.

        You can define your own metric builder as long as it returns a callable that
        accepts the model output spikes and labels.
    """
    def metric_fn(spk_out, y, **kwargs):
        return accuracy_population(
            spk_out,
            y,
            num_classes=metric_cfg.get("num_classes", 2),
        )

    return metric_fn

# note that those below metrics are similar to the accuracy metrics. They follow similar paths. Any metric can be generated if is in function of output spikes and target label, and parameters of those function can even be tuned
# with cfg dict.
def build_precision_metric(metric_cfg):
    def metric_fn(spk_out, y, **kwargs):
        return float(
            precision_population(
                spk_out,
                y,
                num_classes=metric_cfg.get("num_classes", 2),
                positive_class=metric_cfg.get("positive_class", 1),
            )
        )

    return metric_fn


def build_recall_metric(metric_cfg):
    def metric_fn(spk_out, y, **kwargs):
        return float(
            recall_population(
                spk_out,
                y,
                num_classes=metric_cfg.get("num_classes", 2),
                positive_class=metric_cfg.get("positive_class", 1),
            )
        )

    return metric_fn


def build_f1_metric(metric_cfg):
    def metric_fn(spk_out, y, **kwargs):
        return float(
            f1_population(
                spk_out,
                y,
                num_classes=metric_cfg.get("num_classes", 2),
                positive_class=metric_cfg.get("positive_class", 1),
            )
        )

    return metric_fn


def main():
    """
        Configure and run hyperparameter optimization.

        The goal is to maximize validation recall.

        This is useful for binary audio classification tasks where detecting the
        positive class is especially important. For example, in an imbalanced
        detection problem, recall measures how many true positive examples are
        successfully found.
    """


    hpo = HPOWrapper(
        model_builder=build_snn,
        dataset_builder=build_datasets,
        train_fn=train_one_trial,
        loss_builders={"main": build_main_loss},
        metric_builders={
            "acc": build_acc_metric,
            "f1": build_f1_metric,
            "precision": build_precision_metric,
            "recall": build_recall_metric,
        },
        metric="val_recall",
        mode="max",
        experiment_name="snn_tau_search_frontend_fluct_init",
    )


    # ---------------------------------------------------------------------
    # Model hyperparameters
    # ---------------------------------------------------------------------
    # These parameters define the SNN architecture and neuron dynamics.
    # Ray Tune will sample several values and run parallel trials.
    # Later, it will choose promising parameter regions based on previous trials,
    # while maximizing the selected metric.
    
    hpo.bind_model_param("in_ch", 13)
    hpo.bind_model_param("num_classes", 2)
    hpo.bind_model_param("neurons_per_class", 32)
    hpo.bind_model_param("dt", 2e-3)
    hpo.bind_model_param("tau_l1", tune.loguniform(5e-3, 20e-3))
    hpo.bind_model_param("tau_l2", tune.loguniform(10e-3, 50e-3))
    hpo.bind_model_param("tau_l3", tune.loguniform(20e-3, 90e-3))
    hpo.bind_model_param("tau_l4", tune.loguniform(40e-3, 120e-3))
    hpo.bind_model_param("device", "gpu")

    hpo.bind_model_param("init_train_loader_path", "/home/workspaces/autoresearch/data/trainb32plus.pt")
    hpo.bind_model_param("init_verbose", False)


    # ---------------------------------------------------------------------
    # Initialization parameters
    # ---------------------------------------------------------------------
    # The frontend firing-rate initialization encourages the first layer to
    # produce a reasonable spike rate before training.
    #
    # The fluctuation initialization adjusts the network so that internal
    # activity starts in a useful dynamic range.

    hpo.bind_model_param("use_frontend_firing_init", True)
    hpo.bind_model_param("frontend_target_fr", 0.20)
    hpo.bind_model_param("frontend_n_batches", 2)
    hpo.bind_model_param("frontend_epsilon", 0.02)

    hpo.bind_model_param("use_fluct_init", True)
    hpo.bind_model_param("fluct_xi_target", tune.uniform(1.0, 3.0))
    hpo.bind_model_param("fluct_alpha", tune.uniform(0.7, 1.0))
    hpo.bind_model_param("fluct_n_batches", tune.choice([1, 2, 4]))

    hpo.bind_optimizer_param("name", tune.choice(["adam", "adamax"]))
    hpo.bind_optimizer_param("lr", tune.loguniform(1e-4, 1e-2))


    # ---------------------------------------------------------------------
    # Loss configuration
    # ---------------------------------------------------------------------
    # The main loss combines:
    #
    # 1. A population loss that encourages the correct class population to
    #    spike at the desired rate.
    #
    # 2. A firing-rate regularizer that keeps activity in hidden layers under
    #    control and prevents excessive spiking.

    # Those are defined as CompositeLoss previously

    hpo.bind_loss_param("main", "correct_rate", tune.uniform(0.25, 0.45))
    hpo.bind_loss_param("main", "incorrect_rate", tune.uniform(0.01, 0.10))
    hpo.bind_loss_param("main", "multipliers_fr", tune.choice([[1, 2, 3, 4], [10, 20, 5, 100]]))
    hpo.bind_loss_param("main", "offsets_fr", tune.choice([[0.1, 0.2, 0.3, 0.4], [0.01, 0.05, 0.3, 0.3]]))


    # ---------------------------------------------------------------------
    # Metrics
    # ---------------------------------------------------------------------
    # Since this is a binary classification task, we track:
    #
    # - accuracy
    # - precision
    # - recall
    # - F1 score
    #
    # The optimization target is validation recall. Note that all these metrics are as well defined previously, hence you are free to define your own metric in function of spikes and target label.
    hpo.bind_metric_param("acc", "num_classes", tune.choice([2]))
    hpo.bind_metric_param("f1", "num_classes", tune.choice([2]))
    hpo.bind_metric_param("f1", "positive_class", tune.choice([1]))
    hpo.bind_metric_param("precision", "num_classes", tune.choice([2]))
    hpo.bind_metric_param("precision", "positive_class", tune.choice([1]))
    hpo.bind_metric_param("recall", "num_classes", tune.choice([2]))
    hpo.bind_metric_param("recall", "positive_class", tune.choice([1]))


    # ---------------------------------------------------------------------
    # Training schedule and early stopping
    # ---------------------------------------------------------------------
    # Different learning-rate schedulers are tested during HPO.
    # Early stopping ends a trial once validation recall reaches the target.
    hpo.bind_scheduler_param(
        "cfg",
        tune.choice(
            [
                {"name": "none"},
                {
                    "name": "plateau",
                    "mode": "min",
                    "factor": 0.3,
                    "patience": 8,
                    "threshold": 1e-3,
                    "cooldown": 1,
                    "min_lr": 1e-6,
                },
                {
                    "name": "step",
                    "step_size": 10,
                    "gamma": 0.5,
                },
                {
                    "name": "cosine",
                    "T_max": 60,
                    "eta_min": 1e-6,
                },
                {
                    "name": "multistep",
                    "milestones": [20, 40],
                    "gamma": 0.3,
                },
            ]
        ),
    )

    hpo.bind_early_stop_param("policy", "threshold")
    hpo.bind_early_stop_param("metric", "val_recall")
    hpo.bind_early_stop_param("mode", "max")
    hpo.bind_early_stop_param("value", 0.93)
    hpo.bind_early_stop_param("min_epoch", 10)

    hpo.bind_nwave_scheduler_param("surrogate.policy", tune.choice(["linear", "cosine"]))
    hpo.bind_nwave_scheduler_param("surrogate.start_value", tune.choice([10.0, 15.0]))
    hpo.bind_nwave_scheduler_param("surrogate.end_value", tune.choice([25.0, 30.0]))
    hpo.bind_nwave_scheduler_param("sign_annealing.alpha_start", tune.choice([0.5, 1.0]))
    hpo.bind_nwave_scheduler_param("sign_annealing.alpha_end", tune.choice([10.0, 20.0]))

    hpo.bind_trainer_param("epochs", 60)
    hpo.bind_trainer_param("device", "cuda")

    hpo.bind_data_param("train_loader_file", "trainb32plus.pt")
    hpo.bind_data_param("val_loader_file", "valb32plus.pt")

    hpo.set_resources(cpu=2, gpu=0.25)
    hpo.set_max_concurrent_trials(3)
    hpo.set_num_samples(20)


    # ---------------------------------------------------------------------
    # Run the search
    # ---------------------------------------------------------------------
    # The dataset folder contains the serialized training and validation
    # DataLoaders for the audio binary classification task.
    results = hpo.run("/home/workspaces/autoresearch/data")
    print(results)
    print(hpo.best_config())


if __name__ == "__main__":
    main()
