# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/00_core.ipynb (unless otherwise specified).

__all__ = ['download_dataset', 'make_point_clouds', 'load_shapes', 'make_gravitational_waves', 'display_large',
           'rf_feature_importance', 'plot_feature_importance', 'plot_regression_tree', 'plot_predictions']

# Cell
import pandas as pd
from nbdev.showdoc import *
import os
import gdown
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.tree import plot_tree
from pathlib import Path
from typing import List

# Cell
def download_dataset(dataset_name: str):
    """Download datasets from Google Drive."""

    name_to_id = {
        "susy.csv.gz": "1rnR1v-BkMOtzV80R7jFyU1cwO3fGYrQs",
        "susy.feather": "1PxCruwO42GV7FKtwZDXah7iGjDib7YPM",
        "susy_train.feather": "1ezeCZycZ3BrEh-qOLiSJF40YowYEbbTH",
        "susy_test.feather": "1UM8sheb4jzQa16haG6HnVbpJCxZwN2yE",
        "susy_sample.feather": "1l4x_uBeup4eciLDK4YjnfY_G8yTpXLkP",
        "shapes.zip": "1qV-BHzkgwHt8r61Ycf4p9l0CllOQr7I9",
        "Cells.jpg": "1D841Ny9DXAcpgsOTmhgGTahAp4yCWaBh",
        "BlackHole.jpg": "1X1hn31N5n17KQpWw1fXi8Da5CjTTxv0i",
        "diagrams_basic.pkl": "1TJiqCarlcPDk5he_ABscizWhsm_ZIeMG",
        "diagrams.pkl": "1oq9j6oBclSK1u_eVfFeWWktNVuD31neh",
        "gravitational-wave-signals.npy": "1h5R_ceSLw3QsIVOgSQ5pRxwyQykCf6S5"
    }

    path = "../data/"
    os.makedirs(path, exist_ok=True)
    gdrive_path = "https://drive.google.com/uc?id="
    if dataset_name in name_to_id:
        if os.path.exists(path + dataset_name):
            print(
                f"Dataset already exists at '{path + dataset_name}' and is not downloaded again."
            )
            return
        try:
            file_url = gdrive_path + name_to_id[dataset_name]
            gdown.download(file_url, path + dataset_name, quiet=True)
        except Exception as e:
            print("Something went wrong during the download! Try again.")
            raise e
        print(f"Download of {dataset_name} dataset complete.")
    else:
        raise KeyError("File not on Google Drive.")

# Cell
def make_point_clouds(n_samples_per_shape: int, n_points: int, noise: float):
    """Make point clouds for circles, spheres, and tori with random noise.
    """
    circle_point_clouds = [
        np.asarray(
            [
                [np.sin(t) + noise * (np.random.rand(1)[0] - 0.5), np.cos(t) + noise * (np.random.rand(1)[0] - 0.5), 0]
                for t in range((n_points ** 2))
            ]
        )
        for kk in range(n_samples_per_shape)
    ]
    # label circles with 0
    circle_labels = np.zeros(n_samples_per_shape)

    sphere_point_clouds = [
        np.asarray(
            [
                [
                    np.cos(s) * np.cos(t) + noise * (np.random.rand(1)[0] - 0.5),
                    np.cos(s) * np.sin(t) + noise * (np.random.rand(1)[0] - 0.5),
                    np.sin(s) + noise * (np.random.rand(1)[0] - 0.5),
                ]
                for t in range(n_points)
                for s in range(n_points)
            ]
        )
        for kk in range(n_samples_per_shape)
    ]
    # label spheres with 1
    sphere_labels = np.ones(n_samples_per_shape)

    torus_point_clouds = [
        np.asarray(
            [
                [
                    (2 + np.cos(s)) * np.cos(t) + noise * (np.random.rand(1)[0] - 0.5),
                    (2 + np.cos(s)) * np.sin(t) + noise * (np.random.rand(1)[0] - 0.5),
                    np.sin(s) + noise * (np.random.rand(1)[0] - 0.5),
                ]
                for t in range(n_points)
                for s in range(n_points)
            ]
        )
        for kk in range(n_samples_per_shape)
    ]
    # label tori with 2
    torus_labels = 2 * np.ones(n_samples_per_shape)

    point_clouds = np.concatenate((circle_point_clouds, sphere_point_clouds, torus_point_clouds))
    labels = np.concatenate((circle_labels, sphere_labels, torus_labels))

    return point_clouds, labels

# Cell
def load_shapes(path: Path, classes: List, n_points: int) -> pd.DataFrame:
    """Load 3D shapes as a single pandas.DataFrame.
    """
    point_clouds = pd.DataFrame(columns=["x", "y", "z", "label"])

    for c in classes:
        for p in path.rglob("*.pts"):
            if c in p.name:
                # ignore color columns
                df = pd.read_csv(p, names=["x", "y", "z", "r", "g", "b"], usecols=["x", "y", "z"], sep=" ").sample(
                    n_points
                )
                df["label"] = p.stem
                point_clouds = point_clouds.append(df, ignore_index=True)

    return point_clouds

# Cell
def make_gravitational_waves(
    path_to_data: Path,
    n_signals: int = 30,
    downsample_factor: int = 2,
    r_min: float = 0.075,
    r_max: float = 0.65,
    n_snr_values: int = 10,
):
    def padrand(V, n, kr):
        cut = np.random.randint(n)
        rand1 = np.random.randn(cut)
        rand2 = np.random.randn(n - cut)
        out = np.concatenate((rand1 * kr, V, rand2 * kr))
        return out

    Rcoef = np.linspace(r_min, r_max, n_snr_values)
    Npad = 500  # number of padding points on either side of the vector
    gw = np.load(path_to_data / "gravitational-wave-signals.npy")
    Norig = len(gw["data"][0])
    Ndat = len(gw["signal_present"])
    ri = []
    k = 0
    N = int(Norig / downsample_factor)

    ncoeff = []
    Rcoeflist = []

    for j in range(n_signals):
        ncoeff.append(10 ** (-19) * (1 / Rcoef[j % n_snr_values]))
        Rcoeflist.append(Rcoef[j % n_snr_values])

    x = []
    noisy_signals = []
    gw_signals = []
    y = []
    pcloud = []
    maxs = 0.0
    k = 0
    labels = np.zeros(n_signals)

    for j in range(n_signals):
        signal = gw["data"][j % Ndat][range(0, Norig, downsample_factor)]
        sigp = int((np.random.randn() < 0))
        noise = ncoeff[j] * np.random.randn(N)
        labels[j] = sigp
        if sigp == 1:
            rawsig = padrand(signal + noise, Npad, ncoeff[j])
            if k == 0:
                k = 1
        else:
            rawsig = padrand(noise, Npad, ncoeff[j])
        noisy_signals.append(rawsig.copy())
        gw_signals.append(signal)

    return noisy_signals, gw_signals, labels

# Cell
def display_large(df):
    """Displays up to 1000 columns and rows of pandas.DataFrame or pandas.Series objects."""
    with pd.option_context("display.max_rows", 1000, "display.max_columns", 1000):
        display(df)

# Cell
def rf_feature_importance(fitted_model, df):
    "Creates a pandas.Dataframe of a Random Forest's feature importance per column."
    return pd.DataFrame(
        {"Column": df.columns, "Importance": fitted_model.feature_importances_}
    ).sort_values("Importance", ascending=False)

# Cell
def plot_feature_importance(feature_importance):
    fig, ax = plt.subplots(figsize=(12, 8))
    return sns.barplot(y="Column", x="Importance", data=feature_importance, color="b")

# Cell
def plot_regression_tree(fitted_model, feature_names, fontsize=18):
    # we need to specify the background color because of a quirk in sklearn
    fig, ax = plt.subplots(figsize=(30, 10), facecolor='k')
    # generate tree plot
    plot_tree(
        fitted_model,
        filled=True,
        feature_names=feature_names,
        ax=ax,
        fontsize=fontsize,
    )
    plt.show()

# Cell
def plot_predictions(regressors, X, y, axes, label=None, style="r-", data_style="b.", data_label=None):
    x1 = np.linspace(axes[0], axes[1], len(y))
    y_pred = sum(regressor.predict(x1.reshape(-1, 1)) for regressor in regressors)
    plt.plot(X[:, 0], y, data_style, label=data_label)
    plt.plot(x1, y_pred, style, linewidth=2, label=label)
    if label or data_label:
        plt.legend(loc="upper center", fontsize=16)
    plt.axis(axes)
    plt.ylabel("$y$", fontsize=16)
    plt.xlabel("$X$", fontsize=16)