"""Main module."""

import os
from datetime import datetime
import numpy as np
from scipy.optimize import least_squares
import click
import matplotlib.pyplot as plt

from geometric_calibration.reader import (
    read_img_label_file,
    read_projection_hnc,
    read_projection_raw,
)

from geometric_calibration.utils import (
    deg2rad,
    angle2rotm,
    get_grayscale_range,
    create_camera_matrix,
    project_camera_matrix,
    drag_and_drop_bbs,
    search_bbs_centroids,
)

def calibrate_cbct(projection_dir, bbs_3d, sad, sid):
    """Main CBCT Calibration routines.

    :param projection_dir: path to directory containing .raw files
    :type projection_dir: str
    :param bbs_3d: array containing 3D coordinates of BBs
    :type bbs_3d: numpy.array
    :param sad: nominal source to isocenter (A) distance
    :type sad: float
    :param sid: nominal source to image distance
    :type sid: float
    :return: dictionary with calibration results
    :rtype: dict
    """
    # RCS: room coordinate system
    # A: isocenter

    # Read image labels
    labels_file_path = os.path.join(projection_dir, "imgLabels.txt")
    proj_file, angles = read_img_label_file(labels_file_path)

    # Initialize output dictionary
    results = {
        "proj_angles": [],
        "panel_orientation": [],
        "sid": [],
        "sad": [],
        "isocenter": [],
        "source": [],
        "panel": [],
        "img_center": [],
        "err_init": [],
        "err_final": [],
    }

    # Calibrate views
    with click.progressbar(
        iterable=range(len(angles)), fill_char="=", empty_char=" "
    ) as prog_bar:
        for k in prog_bar:
            proj_path = os.path.join(
                projection_dir, proj_file[k]
            )  # path of the current image

            if k == 0:  # no indications other than nominal values
                # Calibrate first view with drag and drop procedure
                view_results = calibrate_projection(
                    proj_path,
                    bbs_3d,
                    sad=sad,
                    sid=sid,
                    angle=angles[k],
                    angle_offset=0,
                    img_dim=[1024, 768],
                    pixel_size=[0.388, 0.388],
                    search_area=7,
                    drag_and_drop=True,
                )
            else:  # if not first iteration
                # initialize geometry (based on previous optimization)
                angle_offset = angles[k] - angles[k - 1]
                image_center = view_results["img_center"]

                # Calibrate other views without drag and drop procedure
                view_results = calibrate_projection(
                    proj_path,
                    bbs_3d,
                    sad=sad,
                    sid=sid,
                    angle=angles[k - 1],
                    angle_offset=angle_offset,
                    img_dim=[1024, 768],
                    pixel_size=[0.388, 0.388],
                    search_area=7,
                    image_center=image_center,
                    drag_and_drop=False,
                )

            # Update output dictionary
            results["proj_angles"].append(view_results["proj_angle"])
            results["panel_orientation"].append(
                view_results["panel_orientation"]
            )
            results["sid"].append(view_results["sid"])
            results["sad"].append(view_results["sad"])
            results["isocenter"].append(view_results["isocenter"])
            results["source"].append(view_results["source"])
            results["panel"].append(view_results["panel"])
            results["img_center"].append(view_results["img_center"])
            results["err_init"].append(view_results["err_init"])
            results["err_final"].append(view_results["err_final"])

    return results

def calibrate_2d(projection_dir, bbs_3d, sad, sid):
    """Main 2D Calibration routines.

    :param projection_dir: path to directory containing .raw files
    :type projection_dir: str
    :param bbs_3d: array containing 3D coordinates of BBs
    :type bbs_3d: numpy.array
    :param sad: nominal source to isocenter (A) distance
    :type sad: float
    :param sid: nominal source to image distance
    :type sid: float
    :return: dictionary with calibration results
    :rtype: dict
    """
    # RCS: room coordinate system
    # A: isocenter

    # Find projection files in the current folder
    proj_file = []
    angles = []

    for f in os.listdir(projection_dir):
        if ("AP" or "RL") and (".raw" or ".hnc") in f:
            proj_file.append(f)
            if "AP" in f:
                angles.append(0)
            elif "RL" in f:
                angles.append(90)

    # Initialize output dictionary
    results = {
        "proj_angles": [],
        "panel_orientation": [],
        "sid": [],
        "sad": [],
        "isocenter": [],
        "source": [],
        "panel": [],
        "img_center": [],
        "err_init": [],
        "err_final": [],
    }

    # Calibrate views
    with click.progressbar(
        iterable=range(len(angles)), fill_char="=", empty_char=" ",
    ) as prog_bar:
        for k in prog_bar:
            proj_path = os.path.join(
                projection_dir, proj_file[k]
            )  # path of the current image

            # Calibrate views with drag and drop procedure
            view_results = calibrate_projection(
                proj_path,
                bbs_3d,
                sad=sad,
                sid=sid,
                angle=angles[k],
                angle_offset=0,
                img_dim=[2048, 1536],
                pixel_size=[0.388, 0.388],
                search_area=14,
                resolution_factor=2,
                drag_and_drop=True,
            )

            # Update output dictionary
            results["proj_angles"].append(view_results["proj_angle"])
            results["panel_orientation"].append(
                view_results["panel_orientation"]
            )
            results["sid"].append(view_results["sid"])
            results["sad"].append(view_results["sad"])
            results["isocenter"].append(view_results["isocenter"])
            results["source"].append(view_results["source"])
            results["panel"].append(view_results["panel"])
            results["img_center"].append(view_results["img_center"])
            results["err_init"].append(view_results["err_init"])
            results["err_final"].append(view_results["err_final"])

    return results

def calibrate_projection(
    projection_file,
    bbs_3d,
    sad,
    sid,
    angle,
    angle_offset=0,
    img_dim=[1024, 768],
    pixel_size=[0.388, 0.388],
    search_area=7,
    resolution_factor=1,
    image_center=None,
    drag_and_drop=True,
):
    """Calibration of a single projection.

    :param projection_file: path to file
    :type projection_file: str
    :param bbs_3d: 3D coordinates of phantom's reference points
    :type bbs_3d: numpy.array
    :param sad: nominal source to isocenter (A) distance
    :type sad: float
    :param sid: nominal source to image distance
    :type sid: float
    :param angle: gantry angle for current projection
    :type angle: float
    :param angle_offset: angle offset for panel, defaults to 0
    :type angle_offset: int, optional
    :param img_dim: image dimensions in pixels, defaults to [1024, 768]
    :type img_dim: list, optional
    :param pixel_size: pixel dimensions in mm, defaults to [0.388, 0.388]
    :type pixel_size: list, optional
    :param search_area: dimension of reference point's searching area, defaults
     to 7
    :type search_area: int, optional
    :param resolution_factor: resolution factor, when mode is "cbct" this
     parameter equals to 1, in 2D mode is 2 (because resolution is doubled),
     defaults to 1
    :type resolution_factor: int, optional
        :param image_center: [description], defaults to None
    :type image_center: [type], optional
    :param image_center: center of image, defaults to None
    :type image_center: list, optional
    :param drag_and_drop: whether or not perform Drag&Drop correction routines,
     typically set to True for first projection. Defaults to True
    :type drag_and_drop: bool, optional
    :raises Exception: if less than 5 BBs centroids are recognized, optimizer
     automatically fails since calibration can't be consider reliable
    :return: dictionary with calibration results for current projection
    :rtype: dict
    """

    results = {}

    if image_center is None:  # in case image_center is not declared
        image_center = [img_dim[1] / 2, img_dim[0] / 2]
    isocenter = [0, 0, 0]

    # panel orientation (from panel to brandis reference - rotation along y)
    panel_orientation = np.array([0, deg2rad(angle), 0]) + np.array(
        [0, deg2rad(angle_offset), 0]
    )

    # Load projection
    if ".raw" in projection_file:
        img = read_projection_raw(projection_file, img_dim)
    elif ".hnc" in projection_file:
        img = read_projection_hnc(projection_file, img_dim)

    # Project points starting from extrinsic and intrinsic parameters
    # generate proj_matrix (extrinsic and intrinsic parameters)
    T = create_camera_matrix(panel_orientation, sid, sad, pixel_size, isocenter)
    # projected coordinates of brandis on panel plane
    r2d = project_camera_matrix(
        bbs_3d, image_center, T, resolution_factor
    )  # 2d coordinates of reference points

    grayscale_range = get_grayscale_range(img)

    if drag_and_drop is True:
        # Overlay reference bbs with projection
        r2d_corrected = drag_and_drop_bbs(
            projection_path=img,
            bbs_projected=r2d,
            grayscale_range=grayscale_range,
        )

    # Starting from the updated coordinates, define a search area around them
    # and identify the bbs as black pixels inside these areas (brandis are used
    # as probes)
    if drag_and_drop is True:
        bbs_centroid = search_bbs_centroids(
            img=img,
            ref_2d=r2d_corrected,
            search_area=search_area,
            dim_img=img_dim,
            grayscale_range=grayscale_range,
        )
    else:
        bbs_centroid = search_bbs_centroids(
            img=img,
            ref_2d=r2d,
            search_area=search_area,
            dim_img=img_dim,
            grayscale_range=grayscale_range,
        )

    # Calibration - non linear data fitting optimization problem
    index = np.where(~np.isnan(bbs_centroid[:, 0]))[0]

    # Estimated BBs
    bbs_estim_init = bbs_centroid[
        ~np.isnan(bbs_centroid).any(axis=1)
    ]  # not consider if out of searching area

    # Real Brandis BBs
    bbs_real_init = bbs_3d[index, :]

    # x0
    parameters = np.append(panel_orientation, image_center).tolist()
    parameters.append(sid)
    parameters.append(sad)

    # Boundaries
    angle_limit = 0.05
    sid_sad_limit = 1
    low_bound = [
        -angle_limit,
        -np.pi,
        -angle_limit,
        0,
        0,
        sid - sid_sad_limit,
        sad - sid_sad_limit,
    ]
    up_bound = [
        angle_limit,
        np.pi,
        angle_limit,
        img_dim[1],
        img_dim[0],
        sid + sid_sad_limit,
        sad + sid_sad_limit,
    ]

    if index.shape[0] > 5:  # at least 5 BBs
        sol = least_squares(
            fun=calibration_cost_function,
            x0=parameters,
            args=(bbs_real_init, pixel_size, bbs_estim_init, isocenter,),
            method="trf",
            bounds=(low_bound, up_bound)
            # verbose=2,
        )

        sol = sol.x  # Solution found

        panel_orientation_new = np.array(sol[:3])  # New panel orientation
        image_center_new = np.array(sol[3:5])  # New center of image
        sid_new = sol[5]
        sad_new = sol[6]
        isocenter_new = isocenter
    else:
        raise Exception("Cannot properly process last projection. Please Retry")

    # project based on calibration - use new panel orientation,
    # tube and panel position
    T = create_camera_matrix(
        panel_orientation_new, sid_new, sad_new, pixel_size, isocenter_new
    )  # projected coordinates of brandis on panel plane

    bbs_estim_final = project_camera_matrix(
        bbs_3d, image_center_new, T
    )  # projected BBs (considering unknown)

    # calculate improvement
    err_init = bbs_estim_init - r2d[index, :]  # estimated - projected
    err_final = bbs_estim_init - bbs_estim_final[index, :]

    err_init = np.mean(abs(err_init))
    err_final = np.mean(abs(err_final))

    # calculate new source/panel position
    T_new = angle2rotm(
        panel_orientation_new[0],
        panel_orientation_new[1],
        panel_orientation_new[2],
    )
    R_new = T_new[:3, :3]

    source_new = (isocenter_new + (R_new * np.array([0, 0, sad_new])))[:, 2]
    panel_new = (isocenter_new + (R_new * np.array([0, 0, sad_new - sid_new])))[
        :, 2
    ]

    # update with new value
    results["proj_angle"] = angle
    results["panel_orientation"] = panel_orientation_new
    results["sid"] = sid_new
    results["sad"] = sad_new
    results["isocenter"] = isocenter_new
    results["source"] = source_new
    results["panel"] = panel_new
    results["img_center"] = image_center_new
    results["err_init"] = err_init
    results["err_final"] = err_final

    return results

def calibration_cost_function(param, bbs_3d, pixel_size, bbs_2d, isocenter):
    """Cost Function for calibration optimizers.

    :param param: parameters to be optimized
    :type param: list
    :param bbs_3d: 3D coordinates of reference BBs
    :type bbs_3d: numpy.array
    :param pixel_size: pixel dimensions in mm
    :type pixel_size: list
    :param bbs_2d: 2D coordinates of BBs projected on the current image
    :type bbs_2d: numpy.array
    :param isocenter: coordinates of isocenter
    :type isocenter: numpy.array
    :return: cost function value to be minimized
    :rtype: float
    """
    # unknown
    panel_orientation = np.array(param[:3])
    img_center = np.array(param[3:5])
    sid = np.array(param[5])
    sad = np.array(param[6])

    T = create_camera_matrix(
        panel_orientation, sid, sad, pixel_size, isocenter
    )  # projected coordinates of brandis on panel plane
    r2d = project_camera_matrix(
        bbs_3d, img_center, T
    )  # projected bbs (considering unknown)

    delta = r2d - bbs_2d  # Error

    diff = np.square(delta[:, 0]) + np.square(
        delta[:, 1]
    )  # consider both directions

    return diff

def plot_calibration_results(calib_results):
    """Plot source/panel position after calibration.

    :param calib_results: dictionary containing results of a calibration
    :type calib_results: dict
    """
    source_pos = np.array(calib_results["source"])
    panel_pos = np.array(calib_results["panel"])
    isocenter = np.array(calib_results["isocenter"])

    def on_key_pressed(event):
        if event.key == "enter":
            plt.close()

    # Plot panel and source positions (trajectory)
    fig = plt.figure(num="Source/Panel Position")
    fig.canvas.mpl_connect("key_press_event", on_key_pressed)

    ax = fig.add_subplot(111, projection="3d")

    ax.scatter(
        source_pos[:, 0],
        source_pos[:, 1],
        source_pos[:, 2],
        marker=".",
        c="g",
        label="Source Position",
    )

    ax.scatter(
        panel_pos[:, 0],
        panel_pos[:, 1],
        panel_pos[:, 2],
        marker=".",
        c="r",
        label="Panel Position",
    )

    ax.scatter(
        isocenter[0, 0],
        isocenter[0, 1],
        isocenter[0, 2],
        marker=".",
        c="b",
        label="Isocenter Position",
    )

    plt.title("Panel/Source position after calibration\nPress Enter to close")
    ax.set_xlabel("X Label [mm]")
    ax.set_ylabel("Y Label [mm]")
    ax.set_zlabel("Z Label [mm]")
    fig.legend(loc="lower right")
    plt.show()

def save_lut(path, calib_results, mode):
    """Save LUT file for a calibration.

    :param path: path to .raw file directory, where LUT will be saved
    :type path: str
    :param calib_results: dictionary containing results for a calibration
    :type calib_results: dict
    :param calib_results: acquisition modality for calibration
    :type calib_results: string
    """
    angles = calib_results["proj_angles"]
    panel_orientation = calib_results["panel_orientation"]
    image_center = calib_results["img_center"]
    sid = calib_results["sid"]
    sad = calib_results["sad"]

    clock = datetime.now()

    if mode == "cbct":
        filename = "CBCT_LUT_{}_{:02}_{:02}-{:02}_{:02}.txt".format(
            clock.year, clock.month, clock.day, clock.hour, clock.minute,
        )
    elif mode == "2d":
        filename = "2D_LUT_{}_{:02}_{:02}-{:02}_{:02}.txt".format(
            clock.year, clock.month, clock.day, clock.hour, clock.minute,
        )
    output_file = os.path.join(path, filename)

    with open(output_file, "w") as res_file:
        res_file.write(f"#Look Up Table for {mode.upper()} reconstruction\n")
        res_file.write(
            "#Angle (deg) | Panel Orientation(rad) [X  Y  Z] | Image_center(pixel) X Y | SID(mm) | SAD(mm)\n"
        )
        res_file.write(
            "#Date:{}_{}_{}_Time:{}_{}_{}.{}\n".format(
                clock.year,
                clock.month,
                clock.day,
                clock.hour,
                clock.minute,
                clock.second,
                clock.microsecond,
            )
        )
        res_file.write("#\n")
        res_file.write("# --> END OF HEADER. FIXED SIZE: 5 lines. \n")

        for k in range(len(angles)):
            res_file.write(
                "{:6.12f} {:6.12f} {:6.12f} {:6.12f} {:6.12f} {:6.12f} {:6.12f} {:6.12f}\n".format(
                    angles[k],
                    panel_orientation[k][0],
                    panel_orientation[k][1],
                    panel_orientation[k][2],
                    image_center[k][0],
                    image_center[k][1],
                    sid[k],
                    sad[k],
                )
            )

        res_file.write(
            r"# END OF FILE. REQUIRED TO ENSURE '\n' at the end of last calibration line. NO MORE LINES AFTER THIS!!!"
        )


import torchvision
from torchvision import transforms
from torch.utils import data

def load_data_fashion_mnist(batch_size, resize=None, workers=4):
    """
        自动下载FashionMNIST数据集，并返回数据加载器
    """
    trans = [transforms.ToTensor()]
    if resize:
        trans.insert(0, transforms.Resize(resize))
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(
        root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)
    return (data.DataLoader(mnist_train, batch_size, shuffle=True,
                            num_workers=workers),
            data.DataLoader(mnist_test, batch_size, shuffle=False,
                            num_workers=workers))

def get_fashion_mnist_labels(labels):
    """返回Fashion-MNIST数据集的文本标签
    Defined in :numref:`sec_fashion_mnist`"""
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels]


from typing import List, Tuple
import pathlib
import copy

import numpy as np
import networkx as nx
from tqdm.auto import tqdm
from defdap import ebsd
from defdap.quat import Quat

from beta_reconstruction.crystal_relations import (
    unq_hex_syms, hex_syms, unq_cub_syms, burg_trans
)

def calc_beta_oris(alpha_ori: Quat) -> List[Quat]:
    """Calculate the possible beta orientations for a given alpha orientation.

    Uses the Burgers relation and crystal symmetries to calculate beta orientations.

    Parameters
    ----------
    alpha_ori
        Orientation of an alpha grain

    Returns
    -------
    list of Quat
        List of possible beta orientations
    """
    beta_oris = []

    for sym in unq_hex_syms:
        beta_oris.append(burg_trans * sym.conjugate * alpha_ori)

    return beta_oris

def beta_oris_from_cub_sym(
    alpha_ori: Quat,
    unq_cub_sym_idx: int,
    hex_sym_idx: int
) -> List[Quat]:
    """

    Parameters
    ----------
    alpha_ori
        The orientation of the grain in the alpha phase.
    unq_cub_sym_idx

    hex_sym_idx

    Returns
    -------
    list of Quat
        Possible beta orientations from given symmetries

    """
    if not (0 <= unq_cub_sym_idx <= 11):
        raise ValueError("unq_cub_sym_idx must be between 0 and 11 inclusive")
    if not (0 <= hex_sym_idx <= 11):
        raise ValueError("hex_sym_idx must be between 0 and 11 inclusive")

    beta_oris = []

    beta_ori_base = hex_syms[hex_sym_idx].conjugate * alpha_ori

    # all cases have one possible beta orientation
    beta_oris.append(burg_trans * beta_ori_base)

    if unq_cub_sym_idx == 9:
        # two extra possible beta orientations
        # B - unq_hex_syms[1] is C^+_3z:
        beta_oris.append(
            burg_trans * unq_hex_syms[1].conjugate * beta_ori_base
        )
        # C - unq_hex_syms[2] is C^+_6z:
        beta_oris.append(
            burg_trans * unq_hex_syms[2].conjugate * beta_ori_base
        )

    if unq_cub_sym_idx > 9:
        # one extra possible beta orientations
        # D - unq_hex_syms[4] is C'_22:
        beta_oris.append(
            burg_trans * unq_hex_syms[4].conjugate * beta_ori_base
        )

    return beta_oris

def calc_misori_of_variants(
    alpha_ori_inv: Quat,
    neighbour_ori: Quat,
    unq_cub_sym_comps: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
    """Calculate possible symmetry variants between two orientations.

    Calculate all possible sym variants for misorientation between two
    orientations undergoing a Burgers type transformation. Then calculate
    the misorientation to the nearest cubic symmetry, this is the deviation
    to a perfect Burgers transformation.

    Parameters
    ----------
    alpha_ori_inv
        Inverse of first orientation
    neighbour_ori
        Second orientation
    unq_cub_sym_comps
        Components of the unique cubic symmetries

    Returns
    -------
    min_misoris : np.ndarray
       The minimum misorientation for each of the possible beta variants - shape (12, 12)

    min_cub_sym_idx : np.ndarray
       The minimum cubic symmetry index for each of the possible variants - shape (12, 12)

    """
    # calculate all possible S^B_m (eqn 11. from [1]) from the
    # measured misorientation from 2 neighbour alpha grains
    # for each S^B_m calculate the 'closest' cubic symmetry
    # (from reduced subset) and the deviation from this symmetry

    # Vectorised calculation of:
    # hex_sym[j].inv * ((neighbour_ori * alpha_ori_inv) * hex_sym[i])
    # labelled: d = h2.inv * (c * h1)
    hex_sym_comps = Quat.extract_quat_comps(hex_syms)
    c = (neighbour_ori * alpha_ori_inv).quatCoef
    h1 = np.repeat(hex_sym_comps, 12, axis=1)  # outer loop
    h2 = np.tile(hex_sym_comps, (1, 12))  # inner loop
    d = np.zeros_like(h1)

    c_dot_h1 = c[1]*h1[1] + c[2]*h1[2] + c[3]*h1[3]
    c_dot_h2 = c[1]*h2[1] + c[2]*h2[2] + c[3]*h2[3]
    h1_dot_h2 = h1[1]*h2[1] + h1[2]*h2[2] + h1[3]*h2[3]

    d[0] = (c[0]*h1[0]*h2[0] - h2[0]*c_dot_h1 +
            c[0]*h1_dot_h2 + h1[0]*c_dot_h2 +
            h2[1] * (c[2]*h1[3] - c[3]*h1[2]) +
            h2[2] * (c[3]*h1[1] - c[1]*h1[3]) +
            h2[3] * (c[1]*h1[2] - c[2]*h1[1]))
    d[1] = (c[0]*h2[0]*h1[1] + h1[0]*h2[0]*c[1] - c[0]*h1[0]*h2[1] +
            c_dot_h1*h2[1] + c_dot_h2*h1[1] - h1_dot_h2*c[1] +
            h2[0] * (c[2]*h1[3] - c[3]*h1[2]) +
            c[0] * (h1[2]*h2[3] - h1[3]*h2[2]) +
            h1[0] * (c[2]*h2[3] - c[3]*h2[2]))
    d[2] = (c[0]*h2[0]*h1[2] + h1[0]*h2[0]*c[2] - c[0]*h1[0]*h2[2] +
            c_dot_h1*h2[2] + c_dot_h2*h1[2] - h1_dot_h2*c[2] +
            h2[0] * (c[3]*h1[1] - c[1]*h1[3]) +
            c[0] * (h1[3]*h2[1] - h1[1]*h2[3]) +
            h1[0] * (c[3]*h2[1] - c[1]*h2[3]))
    d[3] = (c[0]*h2[0]*h1[3] + h1[0]*h2[0]*c[3] - c[0]*h1[0]*h2[3] +
            c_dot_h1*h2[3] + c_dot_h2*h1[3] - h1_dot_h2*c[3] +
            h2[0] * (c[1]*h1[2] - c[2] * h1[1]) +
            c[0] * (h1[1]*h2[2] - h1[2] * h2[1]) +
            h1[0] * (c[1]*h2[2] - c[2] * h2[1]))

    # Vectorised calculation of:
    # burg_trans * (d * burg_trans.inv)
    # labelled: beta_vars = b * (c * b.inv)
    b = burg_trans.quatCoef
    beta_vars = np.zeros_like(h1)

    b_dot_b = b[1]*b[1] + b[2]*b[2] + b[3]*b[3]
    b_dot_d = b[1]*d[1] + b[2]*d[2] + b[3]*d[3]

    beta_vars[0] = d[0] * (b[0]*b[0] + b_dot_b)
    beta_vars[1] = (d[1] * (b[0]*b[0] - b_dot_b) + 2*b_dot_d*b[1] +
                    2*b[0] * (b[2]*d[3] - b[3]*d[2]))
    beta_vars[2] = (d[2] * (b[0]*b[0] - b_dot_b) + 2*b_dot_d*b[2] +
                    2*b[0] * (b[3]*d[1] - b[1]*d[3]))
    beta_vars[3] = (d[3] * (b[0]*b[0] - b_dot_b) + 2*b_dot_d*b[3] +
                    2*b[0] * (b[1]*d[2] - b[2]*d[1]))

    # calculate misorientation to each of the cubic symmetries
    misoris = np.einsum("ij,ik->jk", beta_vars, unq_cub_sym_comps)
    misoris = np.abs(misoris)
    misoris[misoris > 1] = 1.
    misoris = 2 * np.arccos(misoris)

    # find the cubic symmetry with minimum misorientation for each of
    # the beta misorientation variants
    min_cub_sym_idx = np.argmin(misoris, axis=1)
    min_misoris = misoris[np.arange(144), min_cub_sym_idx]
    # reshape to 12 x 12 for each of the hex sym multiplications
    min_cub_sym_idx = min_cub_sym_idx.reshape((12, 12))
    min_misoris = min_misoris.reshape((12, 12))

    return min_misoris, min_cub_sym_idx

def calc_beta_oris_from_misori(
    alpha_ori: Quat,
    neighbour_oris: List[Quat],
    burg_tol: float = 5
) -> Tuple[List[List[Quat]], List[float]]:
    """Calculate the possible beta orientations for a given alpha
    orientation using the misorientation relation to neighbour orientations.

    Parameters
    ----------
    alpha_ori
        A quaternion representing the alpha orientation

    neighbour_oris
        Quaternions representing neighbour grain orientations

    burg_tol
        The threshold misorientation angle to determine neighbour relations

    Returns
    -------
    list of lists of defdap.Quat.quat
        Possible beta orientations, grouped by each neighbour. Any
        neighbour with deviation greater than the tolerance is excluded.
    list of float
        Deviations from perfect Burgers transformation

    """
    burg_tol *= np.pi / 180.
    # This needed to move further up calculation process
    unq_cub_sym_comps = Quat.extract_quat_comps(unq_cub_syms)

    alpha_ori_inv = alpha_ori.conjugate

    beta_oris = []
    beta_devs = []

    for neighbour_ori in neighbour_oris:

        min_misoris, min_cub_sym_idxs = calc_misori_of_variants(
            alpha_ori_inv, neighbour_ori, unq_cub_sym_comps
        )

        # find the hex symmetries (i, j) from give the minimum
        # deviation from the burgers relation for the minimum store:
        # the deviation, the hex symmetries (i, j) and the cubic
        # symmetry if the deviation is over a threshold then set
        # cubic symmetry to -1
        min_misori_idx = np.unravel_index(np.argmin(min_misoris),
                                          min_misoris.shape)
        burg_dev = min_misoris[min_misori_idx]

        if burg_dev < burg_tol:
            beta_oris.append(beta_oris_from_cub_sym(
                alpha_ori, min_cub_sym_idxs[min_misori_idx], int(min_misori_idx[0])
            ))
            beta_devs.append(burg_dev)

    return beta_oris, beta_devs

def calc_beta_oris_from_boundary_misori(
    grain: ebsd.Grain,
    neighbour_network: nx.Graph,
    quat_array: np.ndarray,
    alpha_phase_id : int,
    burg_tol: float = 5
) -> Tuple[List[List[Quat]], List[float], List[Quat]]:
    """Calculate the possible beta orientations for pairs of alpha and
    neighbour orientations using the misorientation relation to neighbour
    orientations.

    Parameters
    ----------
    grain
        The grain currently being reconstructed

    neighbour_network
        A neighbour network mapping grain boundary connectivity

    quat_array
        Array of quaternions, representing the orientations of the pixels of the EBSD map

    burg_tol :
        The threshold misorientation angle to determine neighbour relations

    Returns
    -------
    list of lists of defdap.Quat.quat
        Possible beta orientations, grouped by each neighbour. Any
        neighbour with deviation greater than the tolerance is excluded.
    list of float
        Deviations from perfect Burgers transformation
    list of Quat
        Alpha orientations

    """
    # This needed to move further up calculation process
    unq_cub_sym_comps = Quat.extract_quat_comps(unq_cub_syms)

    beta_oris = []
    beta_devs = []
    alpha_oris = []

    neighbour_grains = neighbour_network.neighbors(grain)
    neighbour_grains = [grain for grain in neighbour_grains
                        if grain.phaseID == alpha_phase_id]
    for neighbour_grain in neighbour_grains:

        bseg = neighbour_network[grain][neighbour_grain]['boundary']
        # check sense of bseg
        if grain is bseg.grain1:
            ipoint = 0
        else:
            ipoint = 1

        for boundary_point_pair in bseg.boundaryPointPairsX:
            point = boundary_point_pair[ipoint]
            alpha_ori = quat_array[point[1], point[0]]

            point = boundary_point_pair[ipoint - 1]
            neighbour_ori = quat_array[point[1], point[0]]

            min_misoris, min_cub_sym_idxs = calc_misori_of_variants(
                alpha_ori.conjugate, neighbour_ori, unq_cub_sym_comps
            )

            # find the hex symmetries (i, j) from give the minimum
            # deviation from the burgers relation for the minimum store:
            # the deviation, the hex symmetries (i, j) and the cubic
            # symmetry if the deviation is over a threshold then set
            # cubic symmetry to -1
            min_misori_idx = np.unravel_index(np.argmin(min_misoris),
                                              min_misoris.shape)
            burg_dev = min_misoris[min_misori_idx]

            if burg_dev < burg_tol / 180 * np.pi:
                beta_oris.append(beta_oris_from_cub_sym(
                    alpha_ori, min_cub_sym_idxs[min_misori_idx], int(min_misori_idx[0])
                ))
                beta_devs.append(burg_dev)
                alpha_oris.append(alpha_ori)

    return beta_oris, beta_devs, alpha_oris

def count_beta_variants(
    beta_oris: List[Quat],
    possible_beta_oris: List[List[Quat]],
    ori_tol: float
) -> np.ndarray:
    """

    Parameters
    ----------
    beta_oris
        Possible beta orientations from burgers relation, there are always 6
    possible_beta_oris
        Possible beta orientations from misorientations between neighbouring grains
    ori_tol
        Tolerance for binning of the orientations into the possible 6
    Returns
    -------
    list of int:
        The variant count for the grain

    """
    if not possible_beta_oris:
        return np.zeros(6, dtype=int)
    # divide 2 because of 2* in misorientation definition
    ori_tol = np.cos(ori_tol / 2 * np.pi / 180.)
    # flatten list of lists
    possible_beta_oris = [item for sublist in possible_beta_oris for item in sublist]

    misoris = np.empty((len(possible_beta_oris), 6))
    for ori_index, ori in enumerate(possible_beta_oris):
        for other_ori_index, other_ori in enumerate(beta_oris):
            misoris[ori_index, other_ori_index] = ori.misOri(other_ori, "cubic")

    # max is actually min because actual misorientation is arccos of this
    max_misoris_idx = np.nanargmax(misoris, axis=1)
    max_misoris = misoris[np.arange(len(possible_beta_oris)), max_misoris_idx]
    variant_count, _ = np.histogram(max_misoris_idx[max_misoris > ori_tol],
                                    range(0, 7))

    return variant_count

def load_map(
    ebsd_path: str,
    min_grain_size: int = 3,
    boundary_tolerance: int = 3,
    use_kuwahara: bool = False,
    kuwahara_tolerance: int = 5
) -> ebsd.Map:
    """Load in EBSD data and do the required prerequisite computations."""

    ebsd_path = pathlib.Path(ebsd_path)
    if ebsd_path.suffix == ".ctf":
        map_type = "OxfordText"
    elif ebsd_path.suffix == ".crc":
        map_type = "OxfordBinary"
    else:
        raise TypeError("Unknown EBSD map type. Can only read .ctf and .crc files.")

    ebsd_map = ebsd.Map(ebsd_path.with_suffix(""), dataType=map_type)
    ebsd_map.buildQuatArray()

    if use_kuwahara:
        ebsd_map.filterData(misOriTol=kuwahara_tolerance)

    ebsd_map.findBoundaries(boundDef=boundary_tolerance)
    ebsd_map.findGrains(minGrainSize=min_grain_size)

    ebsd_map.calcGrainAvOris()

    ebsd_map.buildNeighbourNetwork()

    return ebsd_map

def modal_variant(alpha_grains: List[ebsd.Grain]) -> np.ndarray:
    """Given a map of grains with variant counts, assign the prior beta
    orientation of the grains to the variant with the highest count.

    Parameters
    ----------
    alpha_grains
        Grains to calculate modal variant

    Returns
    -------
    Array of the modal variants

    """
    modal_variants = np.empty(len(alpha_grains), dtype=np.int8)
    for i, grain in enumerate(alpha_grains):
        variant_count = grain.variant_count
        mode_variant = np.where(variant_count == np.max(variant_count))[0]
        if len(mode_variant) == 1:
            mode_variant = mode_variant[0]
        else:
            #  multiple variants with same max
            mode_variant = -1
        modal_variants[i] = mode_variant

    return modal_variants

def assign_beta_variants(
    ebsd_map: ebsd.Map,
    mode: str = "modal",
    alpha_phase_id: int = 0
):
    """Given a map of grains with variant counts, determine the prior
    beta orientation of the grains.

    Parameters
    ----------
    ebsd_map:
        EBSD map to assign the beta variants for.
    mode
        How to perform beta orientation assignment
            'modal': The beta orientation is assigned to the variant
                     with the highest count.
    alpha_phase_id
        Index of the alpha phase in the EBSD map.

    """
    alpha_grains = [grain for grain in ebsd_map
                    if grain.phaseID == alpha_phase_id]
    if mode == "modal":
        assigned_variants = modal_variant(alpha_grains)
    else:
        raise NotImplementedError(f"Mode '{mode}' is not a recognised "
                                  f"way to assign variants.")

    for grain, assigned_variant in zip(alpha_grains, assigned_variants):
        if assigned_variant >= 0:
            parent_beta_ori = grain.beta_oris[assigned_variant]
        else:
            parent_beta_ori = None

        grain.assigned_variant = assigned_variant
        grain.parent_beta_ori = parent_beta_ori

    print("Assignment of beta variants complete.")

def construct_variant_map(
    ebsd_map: ebsd.Map,
    alpha_phase_id: int = 0
) -> np.ndarray:
    alpha_grains = (grain for grain in ebsd_map
                    if grain.phaseID == alpha_phase_id)
    all_lists = ((grain.grainID, grain.assigned_variant) for grain in alpha_grains)
    grain_ids, assigned_variants = zip(*all_lists)

    # points not part of a grain or other phases (-2) and
    # those that were not reconstructed (-1)
    return ebsd_map.grainDataToMapData(
        assigned_variants, grainIds=grain_ids, bg=-2
    )

def construct_beta_quat_array(
    ebsd_map: ebsd.Map,
    alpha_phase_id: int = 0,
    variant_map: np.ndarray = None,
) -> np.ndarray:
    """Construct

    Parameters
    ----------
    ebsd_map:
        EBSD map to assign the beta variants for.
    alpha_phase_id
        Index of the alpha phase in the EBSD map.

    """
    if variant_map is None:
        variant_map = construct_variant_map(ebsd_map, alpha_phase_id)

    transformations = []
    for sym in unq_hex_syms:
        transformations.append(burg_trans * sym.conjugate)
    trans_comps = Quat.extract_quat_comps(transformations)
    trans_comps = trans_comps[:, variant_map[variant_map >= 0]]

    quat_comps = Quat.extract_quat_comps(ebsd_map.quatArray[variant_map >= 0])
    quat_comps_beta = np.empty_like(quat_comps)

    # transformations[variant] * quat
    quat_comps_beta[0, :] = (trans_comps[0, :] * quat_comps[0, :]
                             - trans_comps[1, :] * quat_comps[1, :]
                             - trans_comps[2, :] * quat_comps[2, :]
                             - trans_comps[3, :] * quat_comps[3, :])
    quat_comps_beta[1, :] = (trans_comps[1, :] * quat_comps[0, :]
                             + trans_comps[0, :] * quat_comps[1, :]
                             - trans_comps[3, :] * quat_comps[2, :]
                             + trans_comps[2, :] * quat_comps[3, :])
    quat_comps_beta[2, :] = (trans_comps[2, :] * quat_comps[0, :]
                             + trans_comps[0, :] * quat_comps[2, :]
                             - trans_comps[1, :] * quat_comps[3, :]
                             + trans_comps[3, :] * quat_comps[1, :])
    quat_comps_beta[3, :] = (trans_comps[3, :] * quat_comps[0, :]
                             + trans_comps[0, :] * quat_comps[3, :]
                             - trans_comps[2, :] * quat_comps[1, :]
                             + trans_comps[1, :] * quat_comps[2, :])
    # swap into positive hemisphere if required
    quat_comps_beta[:, quat_comps_beta[0, :] < 0] *= -1

    beta_quat_array = np.empty_like(ebsd_map.quatArray)
    beta_quat_array[variant_map < 0] = Quat(1, 0, 0, 0)
    for i, idx in enumerate(zip(*np.where(variant_map >= 0))):
        beta_quat_array[idx] = Quat(quat_comps_beta[:, i])

    return beta_quat_array

def create_beta_ebsd_map(
    ebsd_map: ebsd.Map,
    mode: str = 'only_beta',
    beta_quat_array: np.ndarray = None,
    variant_map: np.array = None,
    alpha_phase_id: int = 0,
    beta_phase_id: int = 1,
) -> ebsd.Map:
    """

    Parameters
    ----------
    ebsd_map
    mode
        How to copy data from the input map
            'alone': Only include the reconstructed beta
            'append': Append reconstructed beta to present beta phase
            'add': Create a new phase for reconstructed beta
    beta_quat_array
    variant_map
    alpha_phase_id
    beta_phase_id
    """
    if variant_map is None:
        variant_map = construct_variant_map(
            ebsd_map, alpha_phase_id=alpha_phase_id
        )
    if beta_quat_array is None:
        beta_quat_array = construct_beta_quat_array(
            ebsd_map, variant_map=variant_map
        )

    if mode == 'alone':
        # Create map with only the reconstructed beta
        new_phase = copy.copy(ebsd_map.phases[beta_phase_id])
        new_phase.name += " (recon)"
        phases = [new_phase]

        out_phase_array = np.zeros_like(ebsd_map.phaseArray)
        out_phase_array[variant_map >= 0] = 1

        out_quat_array = beta_quat_array

    elif mode == 'append':
        # Append reconstructed beta to original beta phase
        phases = copy.copy(ebsd_map.phases)

        out_phase_array = np.copy(ebsd_map.phaseArray)
        out_phase_array[variant_map >= 0] = beta_phase_id + 1

        out_quat_array = np.where(variant_map >= 0, beta_quat_array,
                                  ebsd_map.quatArray)

    elif mode == 'add':
        # Create addition phase for the reconstructed beta
        phases = copy.copy(ebsd_map.phases)
        new_phase = copy.copy(ebsd_map.phases[beta_phase_id])
        new_phase.name += " (recon)"
        phases.append(new_phase)

        out_phase_array = np.copy(ebsd_map.phaseArray)
        out_phase_array[variant_map >= 0] = ebsd_map.numPhases + 1

        out_quat_array = np.where(variant_map >= 0, beta_quat_array,
                                  ebsd_map.quatArray)

    else:
        raise ValueError(f"Unknown beta map construction mode '{mode}'")

    out_euler_array = np.zeros((3,) + ebsd_map.shape)
    for i in range(ebsd_map.yDim):
        for j in range(ebsd_map.xDim):
            out_euler_array[:, i, j] = out_quat_array[i, j].eulerAngles()

    beta_ebsd_data = {
        'stepSize': ebsd_map.stepSize,
        'phases': phases,
        'phase': out_phase_array,
        'eulerAngle': out_euler_array,
        'bandContrast': ebsd_map.bandContrastArray
    }

    # TODO: Change so quats can be loaded instead of going via Euler angles
    beta_map = ebsd.Map(beta_ebsd_data, dataType="PythonDict")
    beta_map.quatArray = out_quat_array

    return beta_map

def do_reconstruction(
    ebsd_map: ebsd.Map,
    mode: str = 'average',
    burg_tol: float = 5,
    ori_tol: float = 3,
    alpha_phase_id: int = 0,
    beta_phase_id: int = 1
):
    """Apply beta reconstruction to a ebsd map object.

    The reconstructed beta map is stored directly in the ebsd map (this
    should probably change)

    Parameters
    ----------
    ebsd_map:
        EBSD map to apply reconstruction to
    mode
        How to perform reconstruction
            'average': grain average orientations
            'boundary': grain boundary orientations
            'beta': retained beta
    burg_tol
        Maximum deviation from the Burgers relation to allow (degrees)
    ori_tol: float
        Maximum deviation from a beta orientation (degrees)
    alpha_phase_id: int
        Index of the alpha phase in the EBSD map.
    beta_phase_id: int
        Index of the beta phase in the EBSD map.

    """
    # this is the only function that interacts with the ebsd map/grain objects
    alpha_grains = [grain for grain in ebsd_map
                    if grain.phaseID == alpha_phase_id]
    first = True
    for grain in tqdm(alpha_grains):

        beta_oris = calc_beta_oris(grain.refOri)
        variant_count = np.zeros(6, dtype=int)

        if mode == 'boundary':
            if first:
                print("Using boundary mode.")
                first = False
            possible_beta_oris, beta_deviations, alpha_oris = \
                calc_beta_oris_from_boundary_misori(
                    grain, ebsd_map.neighbourNetwork, ebsd_map.quatArray,
                    alpha_phase_id, burg_tol=burg_tol
                )

            for possible_beta_ori, beta_deviation, alpha_ori in zip(
                    possible_beta_oris, beta_deviations, alpha_oris):

                beta_oris_l = calc_beta_oris(alpha_ori)

                variant_count += count_beta_variants(
                    beta_oris_l, [possible_beta_ori], ori_tol
                )

        elif mode == 'beta':
            if first:
                print("Using beta mode.")
                first = False
            neighbour_grains = ebsd_map.neighbourNetwork.neighbors(grain)
            neighbour_oris = [[grain.refOri] for grain in neighbour_grains
                              if grain.phaseID == beta_phase_id]

            possible_beta_oris = neighbour_oris
            beta_deviations = [0.] * len(neighbour_oris)

            variant_count += count_beta_variants(
                beta_oris, possible_beta_oris, ori_tol
            )

        elif mode == 'average':
            if first:
                print("Using average mode.")
                first = False
            neighbour_grains = ebsd_map.neighbourNetwork.neighbors(grain)
            neighbour_oris = [grain.refOri for grain in neighbour_grains
                              if grain.phaseID == alpha_phase_id]

            # determine the possible beta orientations based on misorientation
            # between neighbouring alpha grains
            possible_beta_oris, beta_deviations = calc_beta_oris_from_misori(
                grain.refOri, neighbour_oris, burg_tol=burg_tol
            )

            variant_count += count_beta_variants(
                beta_oris, possible_beta_oris, ori_tol
            )

        else:
            raise ValueError(f"Unknown reconstruction mode '{mode}'")

        # save results in the grain objects
        grain.beta_oris = beta_oris
        grain.possible_beta_oris = possible_beta_oris
        grain.beta_deviations = beta_deviations
        grain.variant_count = variant_count


import math
import torch
import numpy as np
from torch import autograd, optim
from torch.distributions.multivariate_normal import MultivariateNormal
import scipy.optimize
from tqdm import tqdm

def get_schedule(num, rad=4):
    if num == 1:
        return np.array([0.0, 1.0])
    t = np.linspace(-rad, rad, num)
    s = 1.0 / (1.0 + np.exp(-t))
    return (s - np.min(s)) / (np.max(s) - np.min(s))

def differentiable_annealed_importance_sampling(s:torch.Tensor, log_likelihood, log_q, n_steps,
        step_size, partial=False, gamma=0.9, mass_matrix=None, clip_grad=None,
        lrates=None, betas=None, block_grad=False, is_train=True, rng:np.random.RandomState=None):
    """
    s: partcle state: n_particles x d
    """
    n_particles = s.shape[:-1]
    dim = s.shape[-1]

    if n_steps == 0:
        return - log_q(s) + log_likelihood(s), s

    # if lrates is None:
    #     lrates = step_size * torch.ones(n_steps+1, device=s.device)
    if (type(step_size) is torch.Tensor or type(step_size) is np.ndarray) and len(step_size.shape) > 0:
        assert step_size.shape == s.shape[:-1]
        step_size = step_size[..., None]

    if betas is None:
        # betas = torch.linspace(1.0/n_steps, 1.0, n_steps, device=s.device)
        betas = get_schedule(n_steps+1)

    if mass_matrix is None:
        mass_matrix = torch.eye(dim, device=s.device)
    pi_mean = torch.zeros(dim, device=s.device)
    pi = MultivariateNormal(pi_mean, mass_matrix)
    inverse_mass_matrix = torch.inverse(mass_matrix)

    # s.requires_grad = True

    def log_annealed_prob(beta, s: torch.Tensor):
        return (1 - beta) * log_q(s) + beta * log_likelihood(s)

    def grad_log_annealed_prob(beta, s):
        ''' it's important to set create_graph=True '''
        with torch.enable_grad():
            s.requires_grad_()
            grad = autograd.grad(log_annealed_prob(beta, s).sum(), s, create_graph=is_train)[0]
            if clip_grad is not None:
                max_ = torch.prod(torch.tensor(s.shape)) * clip_grad # last dimension of mu_z
                grad = torch.clamp(grad, -max_, max_)
        return grad

    # sample initial momentum
    def pi_sample(n_particles):
        if rng is None:
            eps = pi.sample(n_particles)
        else:
            eps = torch.tensor(rng.multivariate_normal(pi_mean, mass_matrix, n_particles),
                               dtype=s.dtype, device=s.device)
        return eps

    v = pi_sample(n_particles)

    with torch.set_grad_enabled(is_train):

        elbo = - log_q(s)
        for k in range(1, n_steps+1):
            assert not torch.any(s.isnan()), "Current state has nan values"
            elbo = elbo - pi.log_prob(v)

            # leapfrog
            s = s + step_size / 2 * v @ inverse_mass_matrix
            v = v + step_size * grad_log_annealed_prob(betas[k], s)
            s = s + step_size / 2 * v @ inverse_mass_matrix

            elbo = elbo + pi.log_prob(v)

            if partial:
                # partial_refreshment
                v = gamma * v + math.sqrt(1 - math.pow(gamma, 2)) * pi_sample(n_particles)
            else:
                v = pi_sample(n_particles)

        elbo = elbo + log_likelihood(s)

    return elbo, s

def leapfrogs_and_bounds_optlr(s, log_likelihood, log_q, n_steps,
        step_size, partial=False, gamma=0.9, mass_matrix=None):

    init_log_lrates = math.log(step_size) * np.ones(n_steps)
    bounds = scipy.optimize.Bounds(-np.infty, -1.0)
    # log_lrates = torch.tensor(llrates_np, dtype=s.dtype, device=s.device, requires_grad=True)

    def func_fn(log_lrates):
        log_lrates = torch.tensor(log_lrates, dtype=s.dtype, device=s.device, requires_grad=True)
        elbo, _ = differentiable_annealed_importance_sampling(s, log_likelihood, log_q, n_steps,
            step_size, partial, gamma, mass_matrix, lrates=torch.exp(log_lrates))
        return -elbo.sum().data.cpu().numpy()

    def grad_fn(log_lrates):
        log_lrates = torch.tensor(log_lrates, dtype=s.dtype, device=s.device, requires_grad=True)
        elbo, _ = differentiable_annealed_importance_sampling(s, log_likelihood, log_q, n_steps,
            step_size, partial, gamma, mass_matrix, lrates=torch.exp(log_lrates))
        loss = -elbo.sum()
        return autograd.grad(loss, log_lrates)[0].data.cpu().numpy().astype(np.float64)

    res = scipy.optimize.minimize(func_fn, init_log_lrates, jac=grad_fn, bounds=bounds)
    log_lrates = torch.tensor(res.x, dtype=s.dtype, device=s.device)

    return differentiable_annealed_importance_sampling(s, log_likelihood, log_q, n_steps,
        step_size, partial, gamma, mass_matrix, lrates=torch.exp(log_lrates))


from datetime import timedelta

from airflow import DAG
from airflow.models import Param
from airflow.operators.dummy_operator import DummyOperator
from common.operators.gce import (
    StartGCEOperator,
    StopGCEOperator,
    CloneRepositoryGCEOperator,
    SSHGCEOperator,
)
from common.utils import get_airflow_schedule
from airflow.providers.google.cloud.transfers.bigquery_to_gcs import (
    BigQueryToGCSOperator,
)
from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryExecuteQueryOperator,
)
from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator

from common import macros
from common.alerts import task_fail_slack_alert
from common.config import (
    DAG_FOLDER,
    ENV_SHORT_NAME,
    MLFLOW_BUCKET_NAME,
    SLACK_CONN_ID,
    SLACK_CONN_PASSWORD,
    MLFLOW_URL,
    BIGQUERY_RAW_DATASET,
    BIGQUERY_TMP_DATASET,
)

from dependencies.ml.utils import create_algo_training_slack_block
from datetime import datetime

from jobs.ml.constants import IMPORT_TRAINING_SQL_PATH

DATE = "{{ ts_nodash }}"

# Environment variables to export before running commands
dag_config = {
    "STORAGE_PATH": f"gs://{MLFLOW_BUCKET_NAME}/algo_training_{ENV_SHORT_NAME}/algo_training_two_towers_{DATE}",
    "BASE_DIR": "data-gcp/jobs/ml_jobs/algo_training",
    "MODEL_DIR": "two_towers_model",
    "TRAIN_DIR": "/home/airflow/train",
    "EXPERIMENT_NAME": f"algo_training_qpi_v1.1_{ENV_SHORT_NAME}",
}

# Params
train_params = {
    "config_file_name": "user-qpi-features",
    "batch_size": 8192,
    "validation_steps_ratio": 0.1 if ENV_SHORT_NAME == "prod" else 0.4,
    "embedding_size": 64,
    "train_set_size": 0.95 if ENV_SHORT_NAME == "prod" else 0.8,
    "event_day_number": {"prod": 14, "dev": 365, "stg": 14}[ENV_SHORT_NAME],
}
gce_params = {
    "instance_name": f"algo-training-qpi-{ENV_SHORT_NAME}",
    "instance_type": {
        "dev": "n1-standard-2",
        "stg": "n1-standard-8",
        "prod": "n1-standard-32",
    },
}

schedule_dict = {"prod": "0 12 * * 3", "dev": None, "stg": "0 12 * * 2"}

default_args = {
    "start_date": datetime(2023, 3, 17),
    "on_failure_callback": task_fail_slack_alert,
    "retries": 0,
    "retry_delay": timedelta(minutes=2),
}

with DAG(
    "algo_training_two_towers_qpi",
    default_args=default_args,
    description="Custom training job",
    schedule_interval=get_airflow_schedule(schedule_dict[ENV_SHORT_NAME]),
    catchup=False,
    dagrun_timeout=timedelta(minutes=1440),
    user_defined_macros=macros.default,
    template_searchpath=DAG_FOLDER,
    params={
        "branch": Param(
            default="production" if ENV_SHORT_NAME == "prod" else "master",
            type="string",
        ),
        "config_file_name": Param(
            default=train_params["config_file_name"],
            type="string",
        ),
        "batch_size": Param(
            default=str(train_params["batch_size"]),
            type="string",
        ),
        "validation_steps_ratio": Param(
            default=str(train_params["validation_steps_ratio"]),
            type="string",
        ),
        "embedding_size": Param(
            default=str(train_params["embedding_size"]),
            type="string",
        ),
        "train_set_size": Param(
            default=str(train_params["train_set_size"]),
            type="string",
        ),
        "event_day_number": Param(
            default=str(train_params["event_day_number"]),
            type="string",
        ),
        "input_type": Param(
            default="enriched_clicks",
            type="string",
        ),
        "instance_type": Param(
            default=gce_params["instance_type"][ENV_SHORT_NAME],
            type="string",
        ),
        "instance_name": Param(
            default=gce_params["instance_name"]
            + "-"
            + train_params["config_file_name"],
            type="string",
        ),
        "run_name": Param(
            default=train_params["config_file_name"], type=["string", "null"]
        ),
    },
) as dag:
    start = DummyOperator(task_id="start", dag=dag)

    import_tables = {}
    for table in [
        "recommendation_user_features",
        "recommendation_item_features",
        "training_data_enriched_clicks",
    ]:
        import_tables[table] = BigQueryExecuteQueryOperator(
            task_id=f"import_{table}",
            sql=(IMPORT_TRAINING_SQL_PATH / f"{table}.sql").as_posix(),
            write_disposition="WRITE_TRUNCATE",
            use_legacy_sql=False,
            destination_dataset_table=f"{BIGQUERY_RAW_DATASET}.{table}",
            dag=dag,
        )

    # The params.input_type tells the .sql files which table to take as input
    for dataset in ["training", "validation", "test"]:
        task = BigQueryExecuteQueryOperator(
            task_id=f"import_tmp_{dataset}_table",
            sql=(
                IMPORT_TRAINING_SQL_PATH / f"recommendation_{dataset}_data.sql"
            ).as_posix(),
            write_disposition="WRITE_TRUNCATE",
            use_legacy_sql=False,
            destination_dataset_table=f"{BIGQUERY_TMP_DATASET}.{DATE}_recommendation_{dataset}_data",
            dag=dag,
        )
        import_tables[dataset] = task

    gce_instance_start = StartGCEOperator(
        task_id="gce_start_task",
        preemptible=False,
        instance_name="{{ params.instance_name }}",
        instance_type="{{ params.instance_type }}",
        accelerator_types=[{"name": "nvidia-tesla-t4", "count": 1}],
        retries=2,
        labels={"job_type": "ml"},
    )

    fetch_code = CloneRepositoryGCEOperator(
        task_id="fetch_code",
        instance_name="{{ params.instance_name }}",
        python_version="3.10",
        command="{{ params.branch }}",
        retries=2,
    )

    install_dependencies = SSHGCEOperator(
        task_id="install_dependencies",
        instance_name="{{ params.instance_name }}",
        base_dir=dag_config["BASE_DIR"],
        command="pip install -r requirements.txt --user",
        dag=dag,
    )

    store_data = {}
    for split in ["training", "validation", "test"]:
        store_data[split] = BigQueryToGCSOperator(
            task_id=f"store_{split}_data",
            source_project_dataset_table=f"{BIGQUERY_TMP_DATASET}.{DATE}_recommendation_{split}_data",
            destination_cloud_storage_uris=f"{dag_config['STORAGE_PATH']}/raw_recommendation_{split}_data/data-*.parquet",
            export_format="PARQUET",
            dag=dag,
        )

    store_data["bookings"] = BigQueryToGCSOperator(
        task_id=f"store_bookings_data",
        source_project_dataset_table=f"{BIGQUERY_RAW_DATASET}.training_data_bookings",
        destination_cloud_storage_uris=f"{dag_config['STORAGE_PATH']}/bookings/data-*.parquet",
        export_format="PARQUET",
        dag=dag,
    )

    preprocess_data = {}
    for split in ["training", "validation", "test"]:
        preprocess_data[split] = SSHGCEOperator(
            task_id=f"preprocess_{split}",
            instance_name="{{ params.instance_name }}",
            base_dir=dag_config["BASE_DIR"],
            environment=dag_config,
            command=f"PYTHONPATH=. python {dag_config['MODEL_DIR']}/preprocess.py "
            "--config-file-name {{ params.config_file_name }} "
            f"--input-dataframe-file-name raw_recommendation_{split}_data "
            f"--output-dataframe-file-name recommendation_{split}_data",
            dag=dag,
        )

    train = SSHGCEOperator(
        task_id="train",
        instance_name="{{ params.instance_name }}",
        base_dir=dag_config["BASE_DIR"],
        environment=dag_config,
        command=f"PYTHONPATH=. python {dag_config['MODEL_DIR']}/train.py "
        "--config-file-name {{ params.config_file_name }} "
        f"--experiment-name {dag_config['EXPERIMENT_NAME']} "
        "--batch-size {{ params.batch_size }} "
        "--validation-steps-ratio {{ params.validation_steps_ratio }} "
        "--embedding-size {{ params.embedding_size }} "
        "--seed {{ ds_nodash }} "
        "--run-name {{ params.run_name }}",
        dag=dag,
    )

    evaluate = SSHGCEOperator(
        task_id="evaluate",
        instance_name="{{ params.instance_name }}",
        base_dir=dag_config["BASE_DIR"],
        environment=dag_config,
        command=f"PYTHONPATH=. python evaluate.py "
        f"--experiment-name {dag_config['EXPERIMENT_NAME']} "
        "--config-file-name {{ params.config_file_name }}",
        dag=dag,
    )

    gce_instance_stop = StopGCEOperator(
        task_id="gce_stop_task", instance_name="{{ params.instance_name }}"
    )

    send_slack_notif_success = SlackWebhookOperator(
        task_id="send_slack_notif_success",
        http_conn_id=SLACK_CONN_ID,
        webhook_token=SLACK_CONN_PASSWORD,
        blocks=create_algo_training_slack_block(
            dag_config["EXPERIMENT_NAME"], MLFLOW_URL, ENV_SHORT_NAME
        ),
        username=f"Algo trainer robot - {ENV_SHORT_NAME}",
        icon_emoji=":robot_face:",
    )

    (
        start
        >> [
            import_tables["recommendation_user_features"],
            import_tables["recommendation_item_features"],
        ]
        >> import_tables["training_data_enriched_clicks"]
        >> import_tables["training"]
        >> [import_tables["validation"], import_tables["test"]]
        >> gce_instance_start
        >> fetch_code
        >> install_dependencies
        >> [
            store_data["training"],
            store_data["validation"],
            store_data["test"],
            store_data["bookings"],
        ]
        >> preprocess_data["training"]
        >> preprocess_data["validation"]
        >> preprocess_data["test"]
        >> train
        >> evaluate
        >> gce_instance_stop
        >> send_slack_notif_success
    )


import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

T = 200
h = 1e-2
t = np.arange(start=0, stop=T + h, step=h)

bet, gam = 0.15, 1 / 50
# todo: zmienic poziej na randoma
# S_pocz = np.random.uniform(0.7, 1)
S_start = 0.8
I_start = 1 - S_start
R_start = 0
N = S_start + I_start + R_start  # is const

# using odeint
# ---------------------------------------------------------------------------------------------------------------------#
def two_diff_ode_equation(state, t, bet, gam):
    S, I = state
    return [- bet * I * S / N, bet * I * S / N - gam * I]

def one_diff_equation_ode(state, t, bet, gam):
    S = state[0]
    C = I_start - gam / bet * np.log(S_start) + S_start  # C - const
    return [(-bet / N * S * (gam / bet * np.log(S) - S + C))]

def calc_R(S_arr, I_arr):
    R_arr = np.zeros(len(t))
    for i in range(len(R_arr)):
        R_arr[i] = N - S_arr[i] - I_arr[i]
    return R_arr

def calc_I(S_arr):
    C = I_start - gam / bet * np.log(S_start) + S_start  # C - const
    I_arr = np.zeros(len(t))
    for i in range(len(I_arr)):
        I_arr[i] = gam / bet * np.log(S_arr[i]) - S_arr[i] + C
    return I_arr

def two_equation_ode_plot(t, sym, labelt='$t$', labels=['S', 'I', 'R']):
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 4))

    # plot drawing (S, I)
    for i in range(len(labels) - 1):
        ax.plot(t, sym[:, i], label=labels[i])
    # plot drawing (R)
    ax.plot(t, calc_R(sym[:, 0], sym[:, 1]), label=labels[2])

    ax.set_xlabel(labelt, fontsize=14)
    ax.set_ylabel('stan', fontsize=14)
    ax.set_ylim([0, 1])
    ax.legend()
    plt.show()

def one_equation_ode_plot(t, sym, labelt='$t$', labels=['S', 'I', 'R']):
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 4))

    # plot drawing (S)
    ax.plot(t, sym[:, 0], label=labels[0])
    # plot drawing (I)
    I_arr = calc_I(sym[:, 0])
    ax.plot(t, I_arr, label=labels[2])
    # plot drawing (R)
    ax.plot(t, calc_R(sym[:, 0], I_arr), label=labels[2])

    ax.set_xlabel(labelt, fontsize=14)
    ax.set_ylabel('stan', fontsize=14)
    ax.set_ylim([0, 1])
    ax.legend()
    plt.show()

def two_equation_ode_main():
    start_state = S_start, I_start
    sym = odeint(two_diff_ode_equation, start_state, t, args=(bet, gam))
    two_equation_ode_plot(t, sym, labels=['S', 'I', 'R'])

def one_equation_ode_main():
    start_state = S_start
    sym = odeint(one_diff_equation_ode, start_state, t, args=(bet, gam))
    one_equation_ode_plot(t, sym, labels=['S', 'I', 'R'])

# using manual
# ---------------------------------------------------------------------------------------------------------------------#
S = np.zeros(len(t))
S[0] = S_start
I = np.zeros(len(t))
I[0] = I_start
R = np.zeros(len(t))
R[0] = R_start

def two_diff_equation_manual():
    for i in range(t.size - 1):
        S[i + 1] = S[i] + h * (- bet * I[i] * S[i] / N)
        I[i + 1] = I[i] + h * (bet * I[i] * S[i + 1] / N - gam * I[i])
        R[i + 1] = N - S[i + 1] - I[i + 1]

def one_diff_equation_manual():
    C = I_start - gam / bet * np.log(S_start) + S_start  # C - const
    for i in range(t.size - 1):
        S[i + 1] = S[i] + h * (-bet / N * S[i] * (gam / bet * np.log(S[i]) - S[i] + C))
        I[i + 1] = gam / bet * np.log(S[i + 1]) - S[i + 1] + C
        R[i + 1] = N - S[i + 1] - I[i + 1]

def equation_man_plot(t, sirList, labelt='$t$', labels=['S', 'I', 'R']):
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 4))
    # plot drawing (R, S, I)
    for i in range(len(sirList)):
        ax.plot(t, sirList[i], label=labels[i])
    ax.set_xlabel(labelt, fontsize=14)
    ax.set_ylabel('stan', fontsize=14)
    ax.set_ylim([0, 1])
    ax.legend()
    plt.show()

def two_equation_man_main():
    two_diff_equation_manual()
    equation_man_plot(t, [S, I, R], labels=['S', 'I', 'R'])

def one_equation_man_main():
    one_diff_equation_manual()
    equation_man_plot(t, [S, I, R], labels=['S', 'I', 'R'])

if __name__ == "__main__":
    # one_equation_ode_main()
    # one_equation_man_main()
    # two_equation_ode_main()
    two_equation_man_main()
    exit(0)


import numpy as np

import os
import argparse
#import torch
#import torch.nn as nn
from sklearn.linear_model import LinearRegression
import math
import random
from utils import *

'''
def quantize(data,pred,error_bound):
    radius=32768

    diff = data - pred
    quant_index = (int) (abs(diff)/ error_bound) + 1
    #print(quant_index)
    if (quant_index < radius * 2) :
        quant_index =quant_index>> 1
        half_index = quant_index
        quant_index =quant_index<< 1
        #print(quant_index)
        quant_index_shifted=0
        if (diff < 0) :
            quant_index = -quant_index
            quant_index_shifted = radius - half_index
        else :
            quant_index_shifted = radius + half_index

        decompressed_data = pred + quant_index * error_bound
        #print(decompressed_data)
        if abs(decompressed_data - data) > error_bound :
            #print("b")
            return 0,data
        else:
            #print("c")
            data = decompressed_data
            return quant_index_shifted,data

    else:
        #print("a")
        return 0,data
'''
parser = argparse.ArgumentParser()

parser.add_argument('--error','-e',type=float,default=1e-3)
parser.add_argument('--input','-i',type=str)
parser.add_argument('--output','-o',type=str)
parser.add_argument('--quant','-q',type=str,default="ml2_q.dat")
parser.add_argument('--unpred','-u',type=str,default="ml2_u.dat")
parser.add_argument('--max_step','-s',type=int,default=-1)
parser.add_argument('--min_coeff_level','-cl',type=int,default=99)
parser.add_argument('--rate','-r',type=float,default=1.0)
parser.add_argument('--rlist',type=float,default=0.0,nargs="+")
parser.add_argument('--maximum_rate','-m',type=float,default=10.0)
parser.add_argument('--cubic','-c',type=int,default=1)
parser.add_argument('--multidim','-d',type=int,default=1)
parser.add_argument('--lorenzo_fallback_check','-l',type=int,default=-1)
parser.add_argument('--fallback_sample_ratio','-f',type=float,default=0.01)
#parser.add_argument('--level_rate','-lr',type=float,default=1.0)
parser.add_argument('--anchor_rate','-a',type=float,default=0.0)

parser.add_argument('--size_x','-x',type=int,default=1800)
parser.add_argument('--size_y','-y',type=int,default=3600)
#parser.add_argument('--level','-l',type=int,default=2)
#parser.add_argument('--noise','-n',type=bool,default=False)
#parser.add_argument('--intercept','-t',type=bool,default=False)
args = parser.parse_args()

size_x=args.size_x
size_y=args.size_y
array=np.fromfile(args.input,dtype=np.float32).reshape((size_x,size_y))
if args.lorenzo_fallback_check>=0:
    orig_array=np.copy(array)
rng=(np.max(array)-np.min(array))
error_bound=args.error*rng
max_step=args.max_step
rate=args.rate
if args.rlist!=0:
    max_level=int(math.log(max_step,2))
    rate_list=args.rlist
    if isinstance(rate_list,int):
        rate_list=[rate_list]

    while len(rate_list)<max_level:
        rate_list.insert(0,rate_list[0])
else:
    rate_list=None

qs=[]

us=[]
lorenzo_qs=[]
min_coeff_level=args.min_coeff_level
#anchor=args.anchor
if max_step>0:
    max_level=int(math.log(max_step,2))
    anchor_rate=args.anchor_rate
    if anchor_rate>0:
        anchor_eb=error_bound/anchor_rate
        print("Anchor eb:%f" % anchor_eb)

        if max_level>=min_coeff_level:
            reg_xs=[]
            reg_ys=[]
            for x in range(max_step,size_x,max_step):
                for y in range(max_step,size_y,max_step):
                    reg_xs.append(np.array([array[x-max_step][y-max_step],array[x-max_step][y],array[x][y-max_step]],dtype=np.float64))
                    reg_ys.append(array[x][y])
                    res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                    coef=res.coef_
                    ince=res.intercept_

        for x in range(0,size_x,max_step):
            for y in range(0,size_y,max_step):
                orig=array[x][y]
                if x and y and max_level>=min_coeff_level:
                    reg_block=np.array([array[x-max_step][y-max_step],array[x-max_step][y],array[x][y-max_step]],dtype=np.float64)
                    pred=np.dot(reg_block,coef)+ince

                else:
                    f_01=array[x-max_step][y] if x else 0
                    f_10=array[x][y-max_step] if y else 0

                    f_00=array[x-max_step][y-max_step] if x and y else 0

                    pred=f_01+f_10-f_00

                q,decomp=quantize(orig,pred,anchor_eb)
                qs.append(q)
                if q==0:
                    us.append(decomp)
                array[x][y]=decomp
    else:
        anchor_eb=0
else:
    pass#todo,some preparations before level start
#print(len(qs))

last_x=((size_x-1)//max_step)*max_step
last_y=((size_y-1)//max_step)*max_step
step=max_step//2
level=max_level-1
q_start=len(qs)
u_start=len(us)
cumulated_loss=0.0
while step>0:
    cur_qs=[]
    cur_us=[]
    if rate_list!=None:
        cur_eb=error_bound/rate_list[level]
    else:
        cur_eb=error_bound/min(maximum_rate,(rate**level))
    #cur_eb=error_bound/min(args.maximum_rate,(rate**level))
    cur_array=np.copy(array[0:last_x+1:step,0:last_y+1:step])
    cur_size_x,cur_size_y=cur_array.shape
    #print(cur_size_x,cur_size_y)
    print("Level %d started. Current step: %d. Current error_bound: %s." % (level,step,cur_eb))
    best_preds=None#need to copy
    best_absloss=None
    best_qs=[]#need to copy
    best_us=[]#need to copy

    #linear interp
    absloss=0
    selected_algo="none"

    if level>=min_coeff_level:
        reg_xs=[]
        reg_ys=[]
        for x in range(0,cur_size_x,2):
            for y in range(1,cur_size_y,2):
                reg_xs.append(np.array([cur_array[x][y-1],cur_array[x][y+1]],dtype=np.float64))
                reg_ys.append(cur_array[x][y])
                res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                coef=res.coef_
                ince=res.intercept_

    for x in range(0,cur_size_x,2):
        for y in range(1,cur_size_y,2):
            if y==cur_size_y-1:
                continue
            orig=cur_array[x][y]
            if level>=min_coeff_level:
                pred= np.dot( np.array([cur_array[x][y-1],cur_array[x][y+1]]),coef )+ince
            else:
                pred=(cur_array[x][y-1]+cur_array[x][y+1])/2
            absloss+=abs(orig-pred)
            q,decomp=quantize(orig,pred,cur_eb)
            cur_qs.append(q)

            if q==0:
                cur_us.append(decomp)
                #absloss+=abs(decomp)
            cur_array[x][y]=decomp
    if level>=min_coeff_level:
        reg_xs=[]
        reg_ys=[]
        for x in range(1,cur_size_x,2):
            for y in range(0,cur_size_y,2):
                reg_xs.append(np.array([cur_array[x-1][y],cur_array[x+1][y]],dtype=np.float64))
                reg_ys.append(cur_array[x][y])
                res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                coef=res.coef_
                ince=res.intercept_
    for x in range(1,cur_size_x,2):
        for y in range(0,cur_size_y,2):
            if x==cur_size_x-1:
                continue
            orig=cur_array[x][y]
            if level>=min_coeff_level:
                pred= np.dot( np.array([cur_array[x-1][y],cur_array[x+1][y]]),coef )+ince
            else:
                pred=(cur_array[x-1][y]+cur_array[x+1][y])/2
            absloss+=abs(orig-pred)
            q,decomp=quantize(orig,pred,cur_eb)

            cur_qs.append(q)
            if q==0:
                cur_us.append(decomp)
                #absloss+=abs(decomp)
            cur_array[x][y]=decomp
    if level>=min_coeff_level:
        md_reg_xs=[]
        md_reg_ys=[]
        for x in range(1,cur_size_x,2):
            for y in range(1,cur_size_y,2):
                md_reg_xs.append(np.array([cur_array[x-1][y],cur_array[x+1][y],cur_array[x][y-1],cur_array[x][y+1]],dtype=np.float64))
                md_reg_ys.append(cur_array[x][y])
                md_res=LinearRegression(fit_intercept=True).fit(md_reg_xs, md_reg_ys)
                md_coef=md_res.coef_
                md_ince=md_res.intercept_

    for x in range(1,cur_size_x,2):
        for y in range(1,cur_size_y,2):
            if x==cur_size_x-1 or y==cur_size_y-1:
                continue
            orig=cur_array[x][y]
            if level>=min_coeff_level:
                pred=np.dot(np.array([cur_array[x-1][y],cur_array[x+1][y],cur_array[x][y-1],cur_array[x][y+1]]),md_coef)+md_ince
            else:
                pred=(cur_array[x-1][y]+cur_array[x+1][y]+cur_array[x][y-1]+cur_array[x][y+1])/4
            absloss+=abs(orig-pred)
            q,decomp=quantize(orig,pred,cur_eb)

            cur_qs.append(q)
            if q==0:
                cur_us.append(decomp)
                #absloss+=abs(decomp)
            cur_array[x][y]=decomp

    best_preds=np.copy(cur_array)
    best_absloss=absloss
    best_qs=cur_qs.copy()
    best_us=cur_us.copy()
    selected_algo="interp_linear"
    #print(len(cur_qs))
    #cubic interp
    if args.cubic:
        #print("cubic")
        absloss=0
        cur_qs=[]
        cur_us=[]
        cur_array=np.copy(array[0:last_x+1:step,0:last_y+1:step])#reset cur_array
        if level>=min_coeff_level:
            reg_xs=[]
            reg_ys=[]
            for x in range(0,cur_size_x,2):
                for y in range(3,cur_size_y,2):
                    if y+3>=cur_size_y:
                        continue
                    reg_xs.append(np.array([cur_array[x][y-3],cur_array[x][y-1],cur_array[x][y+1],cur_array[x][y+3]],dtype=np.float64))
                    reg_ys.append(cur_array[x][y])
                    res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                    coef=res.coef_
                    ince=res.intercept_
        for x in range(0,cur_size_x,2):
            for y in range(1,cur_size_y,2):
                if y==cur_size_y-1:
                    continue
                orig=cur_array[x][y]
                if y>=3 and y+3<cur_size_y:
                    if level>=min_coeff_level:
                        pred=np.dot(coef,np.array([cur_array[x][y-3],cur_array[x][y-1],cur_array[x][y+1],cur_array[x][y+3]]) )+ince
                    else:
                        pred=(-cur_array[x][y-3]+9*cur_array[x][y-1]+9*cur_array[x][y+1]-cur_array[x][y+3])/16
                else:
                    pred=(cur_array[x][y-1]+cur_array[x][y+1])/2
                absloss+=abs(orig-pred)
                q,decomp=quantize(orig,pred,cur_eb)
                cur_qs.append(q)

                if q==0:
                    cur_us.append(decomp)
                    #absloss+=abs(decomp)
                cur_array[x][y]=decomp
        if level>=min_coeff_level:
            reg_xs=[]
            reg_ys=[]
            for x in range(3,cur_size_x,2):
                for y in range(0,cur_size_y,2):
                    if x+3>=cur_size_x:
                        continue
                    reg_xs.append(np.array([cur_array[x-3][y],cur_array[x-1][y],cur_array[x+1][y],cur_array[x+3][y]],dtype=np.float64))
                    reg_ys.append(cur_array[x][y])
                    res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                    coef=res.coef_
                    ince=res.intercept_
        for x in range(1,cur_size_x,2):
            for y in range(0,cur_size_y,2):
                if x==cur_size_x-1:
                    continue
                orig=cur_array[x][y]
                if x>=3 and x+3<cur_size_x:
                    if level>=min_coeff_level:
                        pred=np.dot(coef,np.array([cur_array[x-3][y],cur_array[x-1][y],cur_array[x+1][y],cur_array[x+3][y]]) )+ince
                    else:
                        pred=(-cur_array[x-3][y]+9*cur_array[x-1][y]+9*cur_array[x+1][y]-cur_array[x+3][y])/16
                else:
                    pred=(cur_array[x-1][y]+cur_array[x+1][y])/2
                absloss+=abs(orig-pred)
                q,decomp=quantize(orig,pred,cur_eb)

                cur_qs.append(q)
                if q==0:
                    cur_us.append(decomp)
                    #absloss+=abs(decomp)
                cur_array[x][y]=decomp
        if level>=min_coeff_level:
            md_reg_xs=[]
            md_reg_ys=[]
            for x in range(1,cur_size_x,2):
                for y in range(1,cur_size_y,2):
                    md_reg_xs.append(np.array([cur_array[x-1][y],cur_array[x+1][y],cur_array[x][y-1],cur_array[x][y+1]],dtype=np.float64))
                    md_reg_ys.append(cur_array[x][y])
                    md_res=LinearRegression(fit_intercept=True).fit(md_reg_xs, md_reg_ys)
                    md_coef=md_res.coef_
                    md_ince=md_res.intercept_

        for x in range(1,cur_size_x,2):
            for y in range(1,cur_size_y,2):
                if x==cur_size_x-1 or y==cur_size_y-1:
                    continue
                orig=cur_array[x][y]
                if level>=min_coeff_level:
                    pred=np.dot(np.array([cur_array[x-1][y],cur_array[x+1][y],cur_array[x][y-1],cur_array[x][y+1]]),md_coef)+md_ince
                else:
                    pred=(cur_array[x-1][y]+cur_array[x+1][y]+cur_array[x][y-1]+cur_array[x][y+1])/4
                absloss+=abs(orig-pred)
                q,decomp=quantize(orig,pred,cur_eb)

                cur_qs.append(q)
                if q==0:
                    cur_us.append(decomp)
                    #absloss+=abs(decomp)
                cur_array[x][y]=decomp
        if absloss<best_absloss:
            selected_algo="interp_cubic"
            best_preds=np.copy(cur_array)
            best_absloss=absloss
            best_qs=cur_qs.copy()
            best_us=cur_us.copy()
    if args.multidim:
        absloss=0
        cur_qs=[]
        cur_us=[]
        cur_array=np.copy(array[0:last_x+1:step,0:last_y+1:step])#reset cur_array
        if level>=min_coeff_level:
            md_reg_xs=[]
            md_reg_ys=[]
            for x in range(1,cur_size_x,2):
                for y in range(1,cur_size_y,2):
                    md_reg_xs.append(np.array([cur_array[x-1][y-1],cur_array[x-1][y+1],cur_array[x+1][y-1],cur_array[x+1][y+1]],dtype=np.float64))
                    md_reg_ys.append(cur_array[x][y])
                    md_res=LinearRegression(fit_intercept=True).fit(md_reg_xs, md_reg_ys)
                    md_coef=md_res.coef_
                    md_ince=md_res.intercept_
        for x in range(1,cur_size_x,2):
            for y in range(1,cur_size_y,2):
                if x==cur_size_x-1 or y==cur_size_y-1:
                    continue
                orig=cur_array[x][y]
                if level>=min_coeff_level:
                    pred=np.dot(np.array([cur_array[x-1][y-1],cur_array[x-1][y+1],cur_array[x+1][y-1],cur_array[x+1][y+1]]),md_coef)+md_ince
                else:
                    pred=(cur_array[x-1][y-1]+cur_array[x-1][y+1]+cur_array[x+1][y-1]+cur_array[x+1][y+1])/4
                absloss+=abs(orig-pred)
                q,decomp=quantize(orig,pred,cur_eb)

                cur_qs.append(q)
                if q==0:
                    cur_us.append(decomp)
                    #absloss+=abs(decomp)
                cur_array[x][y]=decomp
        if level>=min_coeff_level:
            md_reg_xs=[]
            md_reg_ys=[]
            for x in range(0,cur_size_x):
                for y in range(1-(x%2),cur_size_y,2):
                    if x==cur_size_x-1 or y==cur_size_y-1:
                        continue
                    md_reg_xs.append(np.array([cur_array[x][y-1],cur_array[x][y+1],cur_array[x-1][y],cur_array[x+1][y]],dtype=np.float64))
                    md_reg_ys.append(cur_array[x][y])
                    md_res=LinearRegression(fit_intercept=True).fit(md_reg_xs, md_reg_ys)
                    md_coef=md_res.coef_
                    md_ince=md_res.intercept_

        for x in range(0,cur_size_x):
            for y in range(1-(x%2),cur_size_y,2):

                orig=cur_array[x][y]
                if x and y and x!=cur_size_x-1 and y!=cur_size_y-1:
                    if level>=min_coeff_level:
                        pred=np.dot(md_coef,np.array([cur_array[x][y-1],cur_array[x][y+1],cur_array[x-1][y],cur_array[x+1][y]]))+md_ince

                    else:

                        pred=(cur_array[x][y-1]+cur_array[x][y+1]+cur_array[x-1][y]+cur_array[x+1][y])/4
                elif x==0 or x==cur_size_x-1:
                    pred=(cur_array[x][y-1]+cur_array[x][y+1])/2
                else:
                    pred=(cur_array[x-1][y]+cur_array[x+1][y])/2
                absloss+=abs(orig-pred)
                q,decomp=quantize(orig,pred,cur_eb)
                cur_qs.append(q)

                if q==0:
                    cur_us.append(decomp)
                #absloss+=abs(decomp)
                cur_array[x][y]=decomp
        if absloss<best_absloss:
            selected_algo="interp_multidim"
            best_preds=np.copy(cur_array)
            best_absloss=absloss
            best_qs=cur_qs.copy()
            best_us=cur_us.copy()
    #lorenzo
    '''
    cur_array=np.copy(array[0:last_x+1:step,0:last_y+1:step])#reset cur_array
    absloss=0
    cur_qs=[]
    cur_us=[]

    if max_level>=min_coeff_level:
        reg_xs=[]
        reg_ys=[]
        for x in range(cur_size_x):
            for y in range(1-(x%2),cur_size_y,2-(x%2)):
                if not (x and y):
                    continue
                reg_xs.append(np.array([array[x-1][y-1],array[x-1][y],array[x][y-1]],dtype=np.float64))
                reg_ys.append(array[x][y])
                res=LinearRegression(fit_intercept=True).fit(reg_xs, reg_ys)
                coef=res.coef_
                ince=res.intercept_

    for x in range(cur_size_x):
        for y in range(1-(x%2),cur_size_y,2-(x%2)):
            orig=cur_array[x][y]
            if 0:#if x and y and max_level>=min_coeff_level:
                pred=np.dot(coef,np.array([array[x-1][y-1],array[x-1][y],array[x][y-1]]))+ince
            else:

                f_01=cur_array[x-1][y] if x else 0

                f_10=cur_array[x][y-1] if y else 0

                f_00=cur_array[x-1][y-1] if x and y else 0

                pred=f_01+f_10-f_00

            absloss+=abs(orig-pred)
            q,decomp=quantize(orig,pred,cur_eb)
            cur_qs.append(q)

            if q==0:
                cur_us.append(decomp)
                #absloss+=abs(decomp)
            cur_array[x][y]=decomp
    #print(np.max(np.abs(array[0:last_x+1:step,0:last_y+1:step]-cur_array)))
    if absloss<best_absloss:
        best_preds=np.copy(cur_array)
        best_absloss=absloss
        best_qs=cur_qs.copy()
        best_us=cur_us.copy()
        selected_algo="lorenzo"
    '''
    #Lorenzo fallback
    if level<=args.lorenzo_fallback_check:
        absloss=0
        #cur_qs=[]
        #cur_us=[]
        #cur_array=np.copy(array[0:last_x+1:step,0:last_y+1:step])#reset cur_array
        cur_orig_array=orig_array[0:last_x+1:step,0:last_y+1:step]
        total_points=[(x,y) for x in range(cur_orig_array.shape[0]) for y in range(cur_orig_array.shape[1]) if (max_step<=0 or ((x*step)%max_step!=0 and (y*step)%max_step!=0))]
        min_sampled_points=100
        if len(total_points)<min_sampled_points:
            num_sumples=len(total_points)
            sampled_points=total_points
        else:
            num_sumples=max(min_sampled_points,int(len(total_points)*sample_rate) )
            sampled_points=random.sample(total_points,num_sumples)
        for x,y in sampled_points:
            orig=cur_orig_array[x][y]
            f_01=cur_orig_array[x-1][y] if x else 0
            if x and args.max_step>0 and ((x-1)*step)%max_step==0 and (y*step)%max_step==0:
                f_01+=anchor_eb*(2*np.random.rand()-1)
            elif x:
                f_01+=cur_eb*(2*np.random.rand()-1)

            f_10=cur_orig_array[x][y-1] if y else 0
            if y and args.max_step>0 and (x*step)%max_step==0 and ((y-1)*step)%max_step==0:
                f_10+=anchor_eb*(2*np.random.rand()-1)
            elif y:
                f_10+=cur_eb*(2*np.random.rand()-1)

            f_00=cur_orig_array[x-1][y-1] if x and y else 0
            if x and y and args.max_step>0 and ((x-1)*step)%max_step==0 and ((y-1)*step)%max_step==0:
                f_00+=anchor_eb*(2*np.random.rand()-1)
            elif x and y:
                f_00+=cur_eb*(2*np.random.rand()-1)

            pred=f_01+f_10-f_00

            absloss+=abs(orig-pred)
        #print(absloss*len(total_points)/len(sampled_points))
        #print(best_absloss)
        #print(cumulated_loss)
        if absloss*len(total_points)/len(sampled_points)<best_absloss+cumulated_loss:
            selected_algo="lorenzo_fallback"
            best_absloss=0
            best_preds=array[0:last_x+1:step,0:last_y+1:step]
            best_qs=[]
            best_us=[]

            qs=qs[:q_start]
            us=us[:u_start]
            for x in range(cur_size_x):
                for y in range(cur_size_y):

                    if max_step>0 and (x*step)%max_step==0 and (y*step)%max_step==0:
                        #print(x,y)
                        continue
                    orig=best_preds[x][y]
                    f_01=best_preds[x-1][y] if x else 0

                    f_10=best_preds[x][y-1] if y else 0

                    f_00=best_preds[x-1][y-1] if x and y else 0

                    pred=f_01+f_10-f_00

                    best_absloss+=abs(orig-pred)
                    q,decomp=quantize(orig,pred,cur_eb)
                    best_qs.append(q)
                    if q==0:
                        best_us.append(decomp)
                #absloss+=abs(decomp)
                    best_preds[x][y]=decomp

        #print(len(best_qs))

    mean_l1_loss=best_absloss/len(best_qs)
    if selected_algo!="lorenzo_fallback":
        cumulated_loss+=best_absloss
    else:
        cumulated_loss=best_absloss
    #print(np.max(np.abs(array[0:last_x+1:step,0:last_y+1:step]-best_preds)))
    array[0:last_x+1:step,0:last_y+1:step]=best_preds
    #if args.lorenzo_fallback_check:
    #    print(np.max(np.abs(orig_array-array))/rng)
    qs+=best_qs
    us+=best_us
    #print(len(qs))
    print ("Level %d finished. Selected algorithm: %s. Mean prediction abs loss: %f." % (level,selected_algo,mean_l1_loss))
    step=step//2
    level-=1
    #print(len(qs))
    #print(best_absloss)
    #print(cumulated_loss)

def lorenzo_2d(array,x_start,x_end,y_start,y_end):
    for x in range(x_start,x_end):
        for y in range(y_start,y_end):

            orig=array[x][y]

            f_01=array[x-1][y] if x else 0
            f_10=array[x][y-1] if y else 0

            f_00=array[x-1][y-1] if x and y else 0

            pred=f_01+f_10-f_00

            q,decomp=quantize(orig,pred,error_bound)
            lorenzo_qs.append(q)
            if q==0:
                us.append(decomp)
            array[x][y]=decomp
lorenzo_2d(array,0,last_x+1,last_y+1,size_y)
lorenzo_2d(array,last_x+1,size_x,0,size_y)

quants=np.concatenate( (np.array(lorenzo_qs,dtype=np.int32),np.array(qs,dtype=np.int32) ) )
unpreds=np.array(us,dtype=np.float32)
array.tofile(args.output)
quants.tofile(args.quant)
unpreds.tofile(args.unpred)
'''
for x in range(size_x):
    for y in range(size_y):
        if array[x][y]==orig_array[x][y] and x%max_step!=0 and y%max_step!=0:
            print(x,y)
'''


from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QPushButton, QStackedWidget, QToolButton, QComboBox, \
    QTableWidget, QTableWidgetItem, QLabel
from PyQt5.uic import loadUi
from fileDirectory import FileDirectory

class Ui(QMainWindow):
    def __init__(self):
        super(Ui, self).__init__()
        loadUi('../ui/PICK.ui', self)

        self.mainStackedView = self.findChild(QStackedWidget, 'StackView')

        self.OV_TeamConfigButton = self.findChild(QPushButton, 'OV_TeamConfigButton')
        self.OV_EventConfigButton = self.findChild(QPushButton, 'OV_EventConfigButton')
        self.OV_DirectoryConfigButton = self.findChild(QPushButton, 'OV_DirectoryConfigButton')
        self.OV_VectorConfigButton = self.findChild(QPushButton, 'OV_VectorConfigButton')
        self.OV_LogFileConfigButton = self.findChild(QPushButton, 'OV_LogFileConfigButton')
        self.OV_FilterConfigButton = self.findChild(QPushButton, 'OV_FilterConfigButton')
        self.OV_LogEntryConfigButton = self.findChild(QPushButton, 'OV_LogEntryConfigButton')
        self.OV_ExportConfigButton = self.findChild(QPushButton, 'OV_ExportConfigButton')
        self.OV_ChangeConfigButton = self.findChild(QPushButton, 'OV_ChangeConfigButton')
        self.OV_VectorDBConfigButton = self.findChild(QPushButton, 'OV_VectorDBConfigButton')
        self.OV_IconConfigButton = self.findChild(QPushButton, 'OV_IconConfigButton')
        self.OV_GraphBuilderConfigButton = self.findChild(QPushButton, 'OV_GraphBuilderConfigButton')
        self.OV_NodesConfigInTableButton = self.findChild(QPushButton, 'OV_NodesConfigInTableButton')
        self.OV_NodesConfigInGraphButton = self.findChild(QPushButton, 'OV_NodesConfigInGraphButton')
        self.OV_RelationshipConfigButton = self.findChild(QPushButton, 'OV_RelationshipConfigButton')

        self.BlueTeamToolButton = self.findChild(QToolButton, 'BlueTeamToolButton')
        self.RootDirectoryToolButton = self.findChild(QToolButton, 'RootDirectoryToolButton')
        self.RedTeamToolButton = self.findChild(QToolButton, 'RedTeamToolButton')
        self.WhiteTeamToolButton = self.findChild(QToolButton, 'WhiteTeamToolButton')

        self.OV_TeamConfigButton.clicked.connect(lambda: self.btn(0))
        self.OV_EventConfigButton.clicked.connect(lambda: self.btn(1))
        self.OV_DirectoryConfigButton.clicked.connect(lambda: self.btn(2))
        self.OV_VectorConfigButton.clicked.connect(lambda: self.btn(3))
        self.OV_LogFileConfigButton.clicked.connect(lambda: self.btn(4))
        self.OV_FilterConfigButton.clicked.connect(lambda: self.btn(5))
        self.OV_LogEntryConfigButton.clicked.connect(lambda: self.btn(6))
        self.OV_ExportConfigButton.clicked.connect(lambda: self.btn(7))
        self.OV_ChangeConfigButton.clicked.connect(lambda: self.btn(8))
        self.OV_VectorDBConfigButton.clicked.connect(lambda: self.btn(9))
        self.OV_IconConfigButton.clicked.connect(lambda: self.btn(10))
        self.OV_GraphBuilderConfigButton.clicked.connect(lambda: self.btn(11))
        self.OV_NodesConfigInTableButton.clicked.connect(lambda: self.btn(12))
        self.OV_NodesConfigInGraphButton.clicked.connect(lambda: self.btn(13))
        self.OV_RelationshipConfigButton.clicked.connect(lambda: self.btn(14))

        self.BlueTeamToolButton.clicked.connect(lambda: self.btn(15))
        self.RootDirectoryToolButton.clicked.connect(lambda: self.btn(15))
        self.RedTeamToolButton.clicked.connect(lambda: self.btn(15))
        self.WhiteTeamToolButton.clicked.connect(lambda: self.btn(15))

        # Log entry configuration page
        self.LECtable = self.findChild(QTableWidget, 'LEC_LET_TtableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.LECtable.setItem(i, 0, checkbox)
            combo = QComboBox()
            combo.addItems([' ', '1', '2', '3'])
            self.LECtable.setCellWidget(i, 4, combo)
            i += 1

        # log file configuration table checkboxes
        self.LFGtable = self.findChild(QTableWidget, 'LFT_tableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.LFGtable.setItem(i, 5, checkbox)
            i += 1

        # Vector configuration table checkboxes
        self.VCtable = self.findChild(QTableWidget, 'VC_TableView')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.VCtable.setItem(i, 0, checkbox)
            i += 1

        # Relationship config table checkboxes
        self.RCtable = self.findChild(QTableWidget, 'RelationshipConfigTableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.RCtable.setItem(i, 0, checkbox)
            i += 1

        # Pulled vector database table checkboxes
        self.PullVDBtable = self.findChild(QTableWidget, 'VDBC_PulledTableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.PullVDBtable.setItem(i, 0, checkbox)
            i += 1

        # pushed vector database table check boxes
        self.PushVDBtable = self.findChild(QTableWidget, 'VDBC_PushedTableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.PushVDBtable.setItem(i, 0, checkbox)
            i += 1

        # approval database table check boxes and drop down
        self.ADBtable = self.findChild(QTableWidget, 'VDBC_AS_TableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.ADBtable.setItem(i, 0, checkbox)
            combo = QComboBox()
            combo.addItems([' ', '1', '2', '3'])
            self.ADBtable.setCellWidget(i, 7, combo)
            i += 1

        # Node configuration table check boxes and drop down
        self.NCtable = self.findChild(QTableWidget, 'NCITF_NT_tableWidget')
        i = 1
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.NCtable.setItem(i, 0, checkbox)
            i += 1

        i = 1
        while i < 10:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.NCtable.setItem(0, i, checkbox)
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.NCtable.setItem(10, i, checkbox)
            i += 1

        i = 1
        while i < 10:
            j = 6
            while j < 9:
                combo = QComboBox()
                combo.addItems([' ', '1', '2', '3'])
                self.NCtable.setCellWidget(j, i, combo)
                j += 1
            i += 1

        # Node configuration in graphical view checkbox and combobox
        self.NCVtable = self.findChild(QTableWidget, 'NCIGF_TV_tableWidget')
        i = 1
        while i < 10:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.NCVtable.setItem(i, 0, checkbox)
            i += 1

        i = 6
        while i < 9:
            combo = QComboBox()
            combo.addItems([' ', '1', '2', '3'])
            self.NCVtable.setCellWidget(i, 1, combo)
            i += 1

        checkbox = QTableWidgetItem()
        checkbox.setCheckState(Qt.Unchecked)
        self.NCVtable.setItem(0, 1, checkbox)
        self.NCVtable.setItem(10, 1, checkbox)

        # Icon configuration table checkboxes
        self.ICtable = self.findChild(QTableWidget, 'IC_IT_TableWidget')
        i = 0
        while i < 50:
            checkbox = QTableWidgetItem()
            checkbox.setCheckState(Qt.Unchecked)
            self.ICtable.setItem(i, 0, checkbox)
            i += 1

        self.connectButton = self.findChild(QPushButton, 'TeamConfigConnectpushButton')
        self.connectStatus = self.findChild(QLabel, 'ConnectionStatus')
        self.connectButton.clicked.connect(self.connectButtonClicked)

        self.show()

    def connectButtonClicked(self):
        if self.connectButton.text() == 'Disconnect':
            self.connectButton.setText('Connect')
            self.connectStatus.setText('Not Connected')
        else:
            self.connectStatus.setText('Connected')
            self.connectButton.setText('Disconnect')

    def btn(self, index):
        if index < 15:
            self.StackView.setCurrentIndex(index)
        elif index == 15:
            self.window = FileDirectory()
            self.window.show()

if __name__ == "__main__":
    import sys
    application = QApplication(sys.argv)
    window = Ui()
    application.exec_()


from gtts import gTTS
import wikipedia as wiki
import sys

bad_sections = ['See also', 'References', 'External links', 'Gallery']

# removes the sections which wouldn't be useful in these formats
def remove_extra(page):
    new_sections = page.sections
    for x in bad_sections:
        if x in new_sections:
            new_sections.remove(x)
    return new_sections

def to_console(page):
    new_sections = remove_extra(page)
    for x in range(len(new_sections)):
        print(new_sections[x])
        print(page.section(new_sections[x]))

"""
currently has issues encoding, certain characters cause issues.
"""
def to_file(page, user_input):
    fileName = user_input + "_wikipedia.txt"
    file = open(fileName, "w")
    new_sections = remove_extra(page)
    #try:
    for x in range(len(new_sections)):
        file.write(new_sections[x])
        file.write('\n')
        print(page.section(new_sections[x]))
        file.write(page.section(new_sections[x]))
        file.write('\n')
    #except:
        print("Error! something went wrong when writing to file!")

def menu():
    print("For the given summary would you like to:")
    print("1. Copy to a text file (currently broken)")
    print("2. Print to the console")
    print("3. Create an mp3 file")
    user_input = input("Please select the number: ")
    return user_input

def to_mp3(page, user_input):
    new_sections = remove_extra(page)
    clean_string = ""
    for x in range(len(new_sections)):
        clean_string += new_sections[x]
        clean_string += '\n'
        clean_string += page.section(new_sections[x])
    tts = gTTS(text=clean_string, lang='en')
    tts.save(user_input + ".mp3")
    print("file has been turned into an mp3")

# main
user_input = input("Enter the Wikipedia page you want to search!")
page = wiki.page(user_input)
print(wiki.summary(user_input))
menu_selection = menu()
if menu_selection is "1":
    to_file(page, user_input)
    print("File successfully saved to: " + user_input + "_wikipedia.txt")
elif menu_selection == "2":
    to_console(page)
elif menu_selection == "3":
    to_mp3(page, user_input)
else:
    print("Invalid input!")