#!/usr/bin/env python
# coding: utf-8

# Neural Dynamic Occupancy Models for Breeding Bird Survey Data

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm
import itertools
import re
import multiprocessing

import torch
torch.set_default_dtype(torch.float64)
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torch.backends.cudnn as cudnn
import os
import random

# custom modules
import dataset
import utils

# Set device to the gpu if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cudnn.deterministic = True
cudnn.benchmark = True

# Set seed for reproducibility
def set_seed(seed):
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


set_seed(2147483647)


## Reading and parsing the Breeding Bird Survey Data

bbs_train = dataset.BBSData(dataset.bbs.query("group == 'train'"))
bbs_valid = dataset.BBSData(dataset.bbs.query("group == 'validation'"))

# loading one example
ex_sp, ex_gn, ex_fm, ex_or, ex_l1, ex_x, ex_x_p, ex_y = bbs_train[0]

# checking shapes
nx = ex_x.shape[0]
nt = ex_y.shape[0]
nsp = len(dataset.cat_ix["english"].keys())
ngn = len(dataset.cat_ix["genus"].keys())
nfm = len(dataset.cat_ix["family"].keys())
nor = len(dataset.cat_ix["order"].keys())
nx_p = ex_x_p.shape[-1]

print(f"Dimension of x is {nx}")
print(f"Number of time steps is {nt}")
print(f"Number of species is {nsp}")
print(f"Number of genera is {ngn}")
print(f"Number of families is {nfm}")
print(f"Number of orders is {nor}")
print(f"Dimension of detection covariate vector is {nx_p}")


class MultiNet(nn.Module):
    """ Multispecies neural dynamic occupancy model. """

    embed_dim = 2 ** 6
    h_dim = 2 ** 5
    activation = nn.LeakyReLU()

    def __init__(self, num_sp, num_gn, num_fm, num_or, num_l1, nx, nt, nx_p):
        """Initialize a network object. 
        
        Args
        ----
        num_sp (int): number of species
        num_gn (int): number of genera
        num_fm (int): number of families
        num_or (int): number of orders
        num_l1 (int): number of level 1 ecoregions
        nx (int): number of route-level input features
        nt (int): number of timesteps (years)
        nx_p (int): number of detection/survey-level covariates
        """
        # Define embeddings and fully connected layers
        super(MultiNet, self).__init__()
        self.nt = nt
        self.nx_p = nx_p

        # shared layers
        self.l1_embed = nn.Embedding(num_l1, self.embed_dim)
        self.sp_embed = nn.Embedding(num_sp, self.embed_dim)
        self.gn_embed = nn.Embedding(num_gn, self.embed_dim)
        self.fm_embed = nn.Embedding(num_fm, self.embed_dim)
        self.or_embed = nn.Embedding(num_or, self.embed_dim)

        self.to_h = nn.Sequential(
            nn.Linear(self.embed_dim + nx, self.h_dim),
            self.activation,
            nn.Linear(self.h_dim, self.h_dim),
            self.activation,
            nn.Linear(self.h_dim, self.h_dim),
            self.activation,
        )

        # component hidden layers
        self.to_h_psi0 = nn.Sequential(
            nn.Linear(self.h_dim, self.h_dim), self.activation
        )

        self.to_h_t = nn.GRU(
            input_size=self.h_dim,
            hidden_size=self.h_dim * 3,
            num_layers=2,
            batch_first=True,
        )

        # weight layers
        self.to_psi0_weights = self.sp_embed_to_w(self.h_dim)
        self.to_p_weights = self.sp_embed_to_w(self.h_dim + nx_p)
        self.to_phi_weights = self.sp_embed_to_w(self.h_dim)
        self.to_gamma_weights = self.sp_embed_to_w(self.h_dim)

    def sp_embed_to_w(self, dim_out):
        """ Mapping from an embedding to a weight vector. 
        
        Args
        ----
        dim_out (int): the dimension of the weight vector
        """
        layers = nn.Sequential(
            nn.Linear(self.embed_dim * 4, self.h_dim),
            self.activation,
            nn.Linear(self.h_dim, dim_out),
        )
        return layers

    def forward(self, sp, gn, fm, order, l1, x, x_p):
        """ Make a forward pass for the model.
        
        Args
        ----
        sp (tensor): integer indices for species
        gn (tensor): integer indices for genera
        fm (tensor): integer indices for families
        order (tensor): integer indices for orders
        l1 (tensor): integer indices for level 1 ecoregion
        x (tensor): route-level features
        x_p (tensor): survey-level features
        
        Return
        ------
        A dictionary of parameters
        """

        l1_vec = self.l1_embed(l1)
        sp_vec = torch.cat(
            (
                self.sp_embed(sp),
                self.gn_embed(gn),
                self.fm_embed(fm),
                self.or_embed(order),
            ),
            -1,
        )

        h = self.to_h(torch.cat((l1_vec, x), -1))
        h_psi0 = self.to_h_psi0(h)
        h_t, _ = self.to_h_t(h.view(-1, 1, self.h_dim).repeat(1, self.nt, 1))
        h_t = self.activation(h_t.view(-1, self.nt, self.h_dim, 3))
        h_phi = h_t[:, :-1, :, 0]
        h_gamma = h_t[:, :-1, :, 1]
        h_p = h_t[:, :, :, 2]

        phi_weights = self.to_phi_weights(sp_vec)
        phi = torch.sigmoid(
            torch.bmm(h_phi, phi_weights.view(-1, self.h_dim, 1)).view(
                -1, self.nt - 1
            )
        )

        gamma_weights = self.to_gamma_weights(sp_vec)
        gamma = torch.sigmoid(
            torch.bmm(h_gamma, gamma_weights.view(-1, self.h_dim, 1)).view(
                -1, self.nt - 1
            )
        )

        # for each year, concatenate the route embedding with survey condtion data
        p_weights = self.to_p_weights(sp_vec)
        p_t = list()
        for t in range(self.nt):
            # stack hidden state with survey data
            p_t.append(
                torch.bmm(
                    torch.cat((h_p[:, t, :], x_p[:, t, :]), -1).view(
                        -1, 1, self.h_dim + self.nx_p
                    ),
                    p_weights.view(-1, self.h_dim + self.nx_p, 1),
                ).squeeze()
            )
        logit_p = torch.stack(p_t, dim=-1)

        psi0_weights = self.to_psi0_weights(sp_vec)  # (batch, h_dim)
        psi0 = torch.sigmoid(
            # take the row-wise dot product (route vector * species weights)
            torch.bmm(
                h_psi0.view(-1, 1, self.h_dim),
                psi0_weights.view(-1, self.h_dim, 1),
            )
        ).view(-1, 1)

        out_dict = {
            "phi": phi,
            "gamma": gamma,
            "psi0": psi0,
            "logit_p": logit_p,
            "phi_weights": phi_weights,
            "gamma_weights": gamma_weights,
            "psi0_weights": psi0_weights,
            "p_weights": p_weights,
            "h_p": h_p,
            "h_psi0": h_psi0,
            "h_gamma": h_gamma,
            "h_phi": h_phi,
        }
        return out_dict


# Training the neural network

n_epoch = 5
batch_size = 2 ** 11
print(f"Batch size is {batch_size}")

train_loader = DataLoader(
    bbs_train, batch_size=batch_size, shuffle=True, num_workers=multiprocessing.cpu_count()
)
valid_loader = DataLoader(
    bbs_valid, batch_size=batch_size * 4, num_workers=multiprocessing.cpu_count(), shuffle=True
)

net = MultiNet(
    num_sp=nsp,
    num_gn=ngn,
    num_fm=nfm,
    num_or=nor,
    num_l1=len(dataset.cat_ix["L1_KEY"]),
    nx=nx,
    nt=nt,
    nx_p=nx_p,
)
net.to(device)

optimizer = optim.Adam(net.parameters(), weight_decay=1e-3)

train_loss = list()
valid_loss = list()

print("Training the multispecies model...")
for epoch in range(n_epoch):
    print("Starting epoch " + str(epoch + 1) + " of " + str(n_epoch) + "...")
    with torch.enable_grad():
        net.train()
        train_loss.append(
            utils.fit_epoch(
                net, train_loader, training=True, optimizer=optimizer, pb=True
            )
        )
    with torch.no_grad():
        net.eval()
        valid_loss.append(
            utils.fit_epoch(net, valid_loader, training=False, pb=False)
        )
        print(np.mean(valid_loss[-1]))

print("Validation loss for the multispecies model:")
print([np.mean(l) for l in valid_loss])


# Make and save predictions for each species
print("Writing multispecies output for training & validation data.")
for sp in tqdm(dataset.cat_ix["english"].keys()):
    net.eval()
    with torch.no_grad():
        sp_df = dataset.bbs.query("english == @sp & group != 'test'")
        years = list(sp_df.filter(regex="^[0-9]{4}", axis=1))
        sp_ds = dataset.BBSData(sp_df)
        sp_loader = DataLoader(
            sp_ds,
            batch_size=batch_size,
            shuffle=False,
            num_workers=multiprocessing.cpu_count(),
            pin_memory=True,
        )

        sp_p = []
        sp_gamma = []
        sp_phi = []
        sp_psi0 = []
        for i_batch, xy in enumerate(sp_loader):
            _, out = utils.bbs_nll(xy, model=net)
            sp_p.append(torch.sigmoid(out["logit_p"]).cpu().detach().numpy())
            sp_gamma.append(out["gamma"].cpu().detach().numpy())
            sp_phi.append(out["phi"].cpu().detach().numpy())
            sp_psi0.append(out["psi0"].cpu().detach().numpy())

        p_df = pd.DataFrame(
            np.concatenate(sp_p), columns=["p_" + y for y in years]
        )
        gamma_df = pd.DataFrame(
            np.concatenate(sp_gamma), columns=["gamma_" + y for y in years[:-1]]
        )
        phi_df = pd.DataFrame(
            np.concatenate(sp_phi), columns=["phi_" + y for y in years[:-1]]
        )
        psi0_df = pd.DataFrame(np.concatenate(sp_psi0), columns=["psi0"])

        sp_name_lower = re.sub(" ", "_", sp.lower())
        out_path = f"out/{sp_name_lower}_nnet.csv"
        res = pd.concat(
            (
                sp_df.filter(
                    regex="^[0-9]{4}$|sp.bbs|route_id|Lat|Lon", axis=1
                ).reset_index(),
                p_df,
                gamma_df,
                phi_df,
                psi0_df,
            ),
            axis=1,
        )
        res.to_csv(out_path, index=False, na_rep="NA")


# Training species-specific models
class OneSpeciesNet(nn.Module):
    """ Single species neural dynamic occupancy model. """

    embed_dim = 2 ** 3
    h_dim = 2 ** 2
    activation = nn.LeakyReLU()

    def __init__(self, num_l1, nx, nt, nx_p):
        """Initialize a network object. 
        
        Args
        ----
        num_l1 (int): number of level 1 ecoregions
        nx (int): number of route-level input features
        nt (int): number of timesteps (years)
        nx_p (int): number of detection/survey-level covariates
        """

        # Define embeddings and fully connected layers
        super(OneSpeciesNet, self).__init__()
        self.nt = nt
        self.nx_p = nx_p

        # shared layers
        self.l1_embed = nn.Embedding(num_l1, self.embed_dim)

        self.to_h = nn.Sequential(
            nn.Linear(self.embed_dim + nx, self.h_dim), self.activation
        )

        # component output layers
        self.to_psi0 = nn.Linear(self.h_dim, 1)
        self.to_p = nn.Linear(self.h_dim + nx_p, 1)
        self.to_phi = nn.Linear(self.h_dim, 1)
        self.to_gamma = nn.Linear(self.h_dim, 1)

    def forward(self, sp, gn, fm, order, l1, x, x_p):
        """ Make a forward pass for the model.

        Note that while species, genus, family, and order are provided as 
        inputs, these are ignored. (Keeping them allows the re-use of the 
        same dataset definition.)
        
        Args
        ----
        sp (tensor): integer indices for species
        gn (tensor): integer indices for genera
        fm (tensor): integer indices for families
        order (tensor): integer indices for orders
        l1 (tensor): integer indices for level 1 ecoregion
        x (tensor): route-level features
        x_p (tensor): survey-level features
        
        Return
        ------
        A dictionary of parameters
        """

        l1_vec = self.l1_embed(l1)
        h = self.to_h(torch.cat((l1_vec, x), -1))

        phi = torch.sigmoid(self.to_phi(h)).repeat(1, self.nt - 1)
        gamma = torch.sigmoid(self.to_gamma(h)).repeat(1, self.nt - 1)
        psi0 = torch.sigmoid(self.to_psi0(h))

        # for each year, concatenate the route embedding with survey condtion data
        p_t = list()
        for t in range(self.nt):
            # stack hidden state with survey data
            p_t.append(self.to_p(torch.cat((h, x_p[:, t, :]), -1)))
        logit_p = torch.stack(p_t, dim=-1).squeeze()  # batch_size X nt
        out_dict = {
            "phi": phi,
            "gamma": gamma,
            "psi0": psi0,
            "logit_p": logit_p,
        }
        return out_dict


def fit_single_species_model(sp, model):
    """ Fit a single species model. 
    
    This function trains a single species model and writes the 
    results to a csv file.
    
    Args: 
    - sp (string): english species name
    - model (nn.Module): the model to use
    """
    sp_name_lower = re.sub(" ", "_", sp.lower())
    out_path = f"out/{sp_name_lower}_ssnet.csv"
    if os.path.isfile(out_path):
        return f"File already exists for {sp}. Moving on..."

    sp_train = dataset.BBSData(
        dataset.bbs.query("group == 'train' & english == @sp")
    )
    sp_valid = dataset.BBSData(
        dataset.bbs.query("group == 'validation' & english == @sp")
    )
    sp_train_loader = DataLoader(
        sp_train, batch_size=batch_size // 4, shuffle=True, num_workers=multiprocessing.cpu_count()
    )
    sp_valid_loader = DataLoader(
        sp_valid, batch_size=batch_size * 4, num_workers=multiprocessing.cpu_count(), shuffle=True
    )
    sp_net = model(
        num_l1=len(dataset.cat_ix["L1_KEY"]), nx=nx, nt=nt, nx_p=nx_p
    )
    sp_net.to(device)

    optimizer = optim.Adam(sp_net.parameters(), weight_decay=1e-3, lr=0.01)
    max_n_epoch = 100

    train_loss = list()
    valid_loss = list()
    for epoch in range(max_n_epoch):
        with torch.enable_grad():
            sp_net.train()
            train_loss.append(
                utils.fit_epoch(
                    sp_net,
                    sp_train_loader,
                    training=True,
                    optimizer=optimizer,
                    pb=False,
                )
            )
        with torch.no_grad():
            sp_net.eval()
            valid_loss.append(
                utils.fit_epoch(
                    sp_net, sp_valid_loader, training=False, pb=False
                )
            )
            if epoch > 2:
                valid_loss_diff = np.mean(valid_loss[-1]) - np.mean(
                    valid_loss[-2]
                )
                if valid_loss_diff >= -0.2:
                    break

    valid_arr = np.concatenate(valid_loss)
    train_arr = np.concatenate(train_loss)

    sp_df = dataset.bbs.query("english == @sp & group != 'test'")
    years = list(sp_df.filter(regex="^[0-9]{4}", axis=1))
    sp_ds = dataset.BBSData(sp_df)
    sp_loader = DataLoader(
        sp_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=multiprocessing.cpu_count(),
        pin_memory=True,
    )
    sp_p = []
    sp_gamma = []
    sp_phi = []
    sp_psi0 = []
    for i_batch, xy in enumerate(sp_loader):
        _, out = utils.bbs_nll(xy, model=sp_net)
        sp_p.append(torch.sigmoid(out["logit_p"]).cpu().detach().numpy())
        sp_gamma.append(out["gamma"].cpu().detach().numpy())
        sp_phi.append(out["phi"].cpu().detach().numpy())
        sp_psi0.append(out["psi0"].cpu().detach().numpy())

    p_df = pd.DataFrame(np.concatenate(sp_p), columns=["p_" + y for y in years])
    gamma_df = pd.DataFrame(
        np.concatenate(sp_gamma), columns=["gamma_" + y for y in years[:-1]]
    )
    phi_df = pd.DataFrame(
        np.concatenate(sp_phi), columns=["phi_" + y for y in years[:-1]]
    )
    psi0_df = pd.DataFrame(np.concatenate(sp_psi0), columns=["psi0"])

    res = pd.concat(
        (
            sp_df.filter(
                regex="^[0-9]{4}$|sp.bbs|route_id|Lat|Lon", axis=1
            ).reset_index(),
            p_df,
            gamma_df,
            phi_df,
            psi0_df,
        ),
        axis=1,
    )
    res.to_csv(out_path, index=False, na_rep="NA")

    del sp_net
    del sp_train
    del sp_valid
    del sp_train_loader
    del sp_valid_loader


# iterate over species and fit a model for each
print("Training and saving single-species results")
for sp in tqdm(dataset.cat_ix["english"].keys()):
    fit_single_species_model(sp, model=OneSpeciesNet)


# Training the final model
final_net = MultiNet(
    num_sp=nsp,
    num_gn=ngn,
    num_fm=nfm,
    num_or=nor,
    num_l1=len(dataset.cat_ix["L1_KEY"]),
    nx=nx,
    nt=nt,
    nx_p=nx_p,
)
final_net.to(device)
final_optimizer = optim.Adam(final_net.parameters(), weight_decay=1e-3)
bbs_retrain = dataset.BBSData(dataset.bbs.query("group != 'test'"))
retrain_loader = DataLoader(
    bbs_retrain, batch_size=batch_size, shuffle=True, num_workers=multiprocessing.cpu_count()
)
retrain_loss = list()

print("Retraining the final model...")
for epoch in range(n_epoch):
    print("Starting epoch " + str(epoch + 1) + " of " + str(n_epoch) + "...")
    with torch.enable_grad():
        final_net.train()
        retrain_loss.append(
            utils.fit_epoch(
                final_net,
                retrain_loader,
                training=True,
                optimizer=final_optimizer,
                pb=True,
            )
        )

## Save species-specific predictions

weight_dir = os.path.join("out", "weights")
if not os.path.exists(weight_dir):
    os.mkdir(weight_dir)


def melt(a, cols=["row_idx", "t", "h_dim", "value"]):
    """ Convert an nd array to a data frame. """
    res = pd.DataFrame([tuple(list(x) + [val]) for x, val in np.ndenumerate(a)])
    res.columns = cols
    return res


print("Saving output from final model...")
for sp in tqdm(dataset.cat_ix["english"].keys()):
    final_net.eval()
    with torch.no_grad():
        sp_df = dataset.bbs.query("english == @sp")
        years = list(sp_df.filter(regex="^[0-9]{4}", axis=1))
        sp_ds = dataset.BBSData(sp_df)
        sp_loader = DataLoader(
            sp_ds,
            batch_size=batch_size,
            shuffle=False,
            num_workers=multiprocessing.cpu_count(),
            pin_memory=True,
        )

        sp_p = []
        sp_gamma = []
        sp_phi = []
        sp_psi0 = []
        h_p = []
        h_phi = []
        h_gamma = []
        h_psi0 = []

        for i_batch, xy in enumerate(sp_loader):
            _, out = utils.bbs_nll(xy, model=final_net)
            sp_p.append(torch.sigmoid(out["logit_p"]).cpu().detach().numpy())
            sp_gamma.append(out["gamma"].cpu().detach().numpy())
            sp_phi.append(out["phi"].cpu().detach().numpy())
            sp_psi0.append(out["psi0"].cpu().detach().numpy())
            h_p.append(out["h_p"].cpu().detach().numpy())
            h_phi.append(out["h_phi"].cpu().detach().numpy())
            h_gamma.append(out["h_gamma"].cpu().detach().numpy())
            h_psi0.append(out["h_psi0"].cpu().detach().numpy())

        p_df = pd.DataFrame(
            np.concatenate(sp_p), columns=["p_" + y for y in years]
        )
        gamma_df = pd.DataFrame(
            np.concatenate(sp_gamma), columns=["gamma_" + y for y in years[:-1]]
        )
        phi_df = pd.DataFrame(
            np.concatenate(sp_phi), columns=["phi_" + y for y in years[:-1]]
        )
        psi0_df = pd.DataFrame(np.concatenate(sp_psi0), columns=["psi0"])

        sp_name_lower = re.sub(" ", "_", sp.lower())
        out_path = f"out/{sp_name_lower}_finalnet.csv"
        res = pd.concat(
            (
                sp_df.filter(
                    regex="^[0-9]{4}$|sp.bbs|route_id|Lat|Lon", axis=1
                ).reset_index(),
                p_df,
                gamma_df,
                phi_df,
                psi0_df,
            ),
            axis=1,
        )
        res.to_csv(out_path, index=False, na_rep="NA")

        # write out weights, which are constant within species
        # (hence we can extract first row ~~ first sample from the minibatch)
        psi0_weight_df = melt(
            out["psi0_weights"][0, :].cpu().detach().numpy(),
            cols=["h_dim", "value"],
        )
        psi0_weight_df["par"] = "psi0"

        p_weight_df = melt(
            out["p_weights"][0, :].cpu().detach().numpy(),
            cols=["h_dim", "value"],
        )
        p_weight_df["par"] = "p"

        phi_weight_df = melt(
            out["phi_weights"][0, :].cpu().detach().numpy(),
            cols=["h_dim", "value"],
        )
        phi_weight_df["par"] = "phi"

        gamma_weight_df = melt(
            out["gamma_weights"][0, :].cpu().detach().numpy(),
            cols=["h_dim", "value"],
        )
        gamma_weight_df["par"] = "gamma"

        weight_df = pd.concat(
            (phi_weight_df, gamma_weight_df, p_weight_df, psi0_weight_df),
            axis=0,
            sort=True,
        )
        weight_df.to_csv(
            os.path.join(weight_dir, f"{sp_name_lower}_weights.csv"),
            index=False,
            na_rep="NA",
        )


# Write out final route vector representations

h_phi_df = melt(
    np.concatenate(h_phi).reshape((-1, final_net.nt - 1, final_net.h_dim))
)
h_phi_df["par"] = "phi"
h_gamma_df = melt(
    np.concatenate(h_gamma).reshape((-1, final_net.nt - 1, final_net.h_dim))
)
h_gamma_df["par"] = "gamma"
h_p_df = melt(np.concatenate(h_p).reshape((-1, final_net.nt, final_net.h_dim)))
h_p_df["par"] = "p"

h_psi0_df = pd.DataFrame(np.concatenate(h_psi0))
h_psi0_df["par"] = "psi0"
h_psi0_df["row_idx"] = list(range(h_psi0_df.shape[0]))

# reshape to permit concatenation
h_psi0_df = pd.melt(h_psi0_df, id_vars=["row_idx", "par"]).rename(
    columns={"variable": "h_dim"}
)

# concatenate all route vectors into one data frame and save
h_df = pd.concat((h_phi_df, h_gamma_df, h_p_df, h_psi0_df), axis=0, sort=True)
h_df.to_csv("out/route_embeddings.csv", index=False)
