"""
Data preparation script for GNN tracking.
This script processes the TrackML dataset and produces graph data on disk.

Author: Gage DeZoort
Reproduced from https://github.com/GageDeZoort/interaction_network_paper
"""

# System
import os
import sys
import time
import argparse
import logging
import multiprocessing as mp
from functools import partial

sys.path.append("../")

# Externals
import yaml
import numpy as np
import pandas as pd
import trackml.dataset
import time

# Locals
from collections import namedtuple

Graph = namedtuple("Graph", ["x", "edge_attr", "edge_index", "y", "pid", "pt"])


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser("prepare.py")
    add_arg = parser.add_argument
    add_arg("config", nargs="?", default="graph-construction/configs/data-driven.yaml")
    add_arg("--n-workers", type=int, default=1)
    add_arg("--task", type=int, default=0)
    add_arg("--n-tasks", type=int, default=1)
    add_arg("-v", "--verbose", action="store_true")
    add_arg("--show-config", action="store_true")
    add_arg("--interactive", action="store_true")
    add_arg("--start-evtid", type=int, default=0)
    add_arg("--end-evtid", type=int, default=10000)
    return parser.parse_args()


def calc_dphi(phi1, phi2):
    """Computes phi2-phi1 given in range [-pi,pi]"""
    dphi = phi2 - phi1
    dphi[dphi > np.pi] -= 2 * np.pi
    dphi[dphi < -np.pi] += 2 * np.pi
    return dphi


def calc_eta(r, z):
    theta = np.arctan2(r, z)
    return -1.0 * np.log(np.tan(theta / 2.0))


def select_segments(
    hits1,
    hits2,
    phi_slope_max,
    z0_max,
    layer1,
    layer2,
    remove_intersecting_edges=False,
    module_map=[],
):
    """
    Construct a list of selected segments from the pairings
    between hits1 and hits2, filtered with the specified
    phi slope and z0 criteria.

    Returns: pd DataFrame of (index_1, index_2), corresponding to the
    DataFrame hit label-indices in hits1 and hits2, respectively.
    """

    # Start with all possible pairs of hits
    keys = ["evtid", "r", "phi", "z", "module_id"]
    hit_pairs = (
        hits1[keys]
        .reset_index()
        .merge(hits2[keys].reset_index(), on="evtid", suffixes=("_1", "_2"))
    )

    # Compute line through the points
    dphi = calc_dphi(hit_pairs.phi_1, hit_pairs.phi_2)
    dz = hit_pairs.z_2 - hit_pairs.z_1
    dr = hit_pairs.r_2 - hit_pairs.r_1
    eta_1 = calc_eta(hit_pairs.r_1, hit_pairs.z_1)
    eta_2 = calc_eta(hit_pairs.r_2, hit_pairs.z_2)
    deta = eta_2 - eta_1
    dR = np.sqrt(deta**2 + dphi**2)
    phi_slope = dphi / dr
    z0 = hit_pairs.z_1 - hit_pairs.r_1 * dz / dr
    # print("max(z0) = ", np.max(abs(z0)))
    # Apply the intersecting line cut
    intersected_layer = dr.abs() < -1
    mid1 = hit_pairs.module_id_1.values
    mid2 = hit_pairs.module_id_2.values
    if remove_intersecting_edges:

        # Innermost barrel layer --> innermost L,R endcap layers
        if (layer1 == 0) and (layer2 == 11 or layer2 == 4):
            z_coord = 71.56298065185547 * dz / dr + z0
            intersected_layer = np.logical_and(z_coord > -490.975, z_coord < 490.975)
        if (layer1 == 1) and (layer2 == 11 or layer2 == 4):
            z_coord = 115.37811279296875 * dz / dr + z0
            intersected_layer = np.logical_and(z_coord > -490.975, z_coord < 490.975)

    # Filter segments according to criteria
    good_seg_mask = (
        (phi_slope.abs() < phi_slope_max)
        & (z0.abs() < z0_max)
        & (intersected_layer == False)
        & (module_map[mid1, mid2])
    )

    dr = dr[good_seg_mask]
    dphi = dphi[good_seg_mask]
    dz = dz[good_seg_mask]
    dR = dR[good_seg_mask]
    return hit_pairs[["index_1", "index_2"]][good_seg_mask], dr, dphi, dz, dR


def construct_graph(
    hits,
    layer_pairs,
    phi_slope_max,
    z0_max,
    feature_names,
    feature_scale,
    evtid="-1",
    remove_intersecting_edges=False,
    module_maps={},
):
    """Construct one graph (e.g. from one event)"""

    t0 = time.time()

    # Loop over layer pairs and construct segments
    layer_groups = hits.groupby("layer")
    segments = []
    seg_dr, seg_dphi, seg_dz, seg_dR = [], [], [], []
    for (layer1, layer2) in layer_pairs:
        module_map = module_maps[(layer1, layer2)]

        # Find and join all hit pairs
        try:
            hits1 = layer_groups.get_group(layer1)
            hits2 = layer_groups.get_group(layer2)
        # If an event has no hits on a layer, we get a KeyError.
        # In that case we just skip to the next layer pair
        except KeyError as e:
            logging.info("skipping empty layer: %s" % e)
            continue
        # Construct the segments
        selected, dr, dphi, dz, dR = select_segments(
            hits1,
            hits2,
            phi_slope_max,
            z0_max,
            layer1,
            layer2,
            remove_intersecting_edges=remove_intersecting_edges,
            module_map=module_map,
        )

        segments.append(selected)
        seg_dr.append(dr)
        seg_dphi.append(dphi)
        seg_dz.append(dz)
        seg_dR.append(dR)

    # Combine segments from all layer pairs
    segments = pd.concat(segments)
    seg_dr, seg_dphi = pd.concat(seg_dr), pd.concat(seg_dphi)
    seg_dz, seg_dR = pd.concat(seg_dz), pd.concat(seg_dR)

    # Prepare the graph matrices
    n_hits = hits.shape[0]
    n_edges = segments.shape[0]
    X = (hits[feature_names].values / feature_scale).astype(np.float32)
    edge_attr = np.stack(
        (
            seg_dr / feature_scale[0],
            seg_dphi / feature_scale[1],
            seg_dz / feature_scale[2],
            seg_dR,
        )
    )

    y = np.zeros(n_edges, dtype=np.float32)

    # We have the segments' hits given by dataframe label,
    # so we need to translate into positional indices.
    # Use a series to map hit label-index onto positional-index.
    hit_idx = pd.Series(np.arange(n_hits), index=hits.index)
    seg_start = hit_idx.loc[segments.index_1].values
    seg_end = hit_idx.loc[segments.index_2].values
    edge_index = np.stack((seg_start, seg_end))

    # Fill the segment, particle labels
    pid = hits.particle_id
    pt = hits.pt
    unique_pid_map = {
        pid_old: pid_new for pid_new, pid_old in enumerate(np.unique(pid.values))
    }
    pid_mapped = pid.map(unique_pid_map)

    pid1 = hits.particle_id.loc[segments.index_1].values
    pid2 = hits.particle_id.loc[segments.index_2].values
    y[:] = pid1 == pid2

    # Correct for multiple true barrel-endcap segments
    layer1 = hits.layer.loc[segments.index_1].values
    layer2 = hits.layer.loc[segments.index_2].values
    true_layer1 = layer1[y > 0.5]
    true_layer2 = layer2[y > 0.5]
    true_pid1 = pid1[y > 0.5]
    true_pid2 = pid2[y > 0.5]
    true_z1 = hits.z.loc[segments.index_1].values[y > 0.5]
    true_z2 = hits.z.loc[segments.index_2].values[y > 0.5]
    pid_lookup = {}
    for p, pid in enumerate(np.unique(true_pid1)):
        pid_lookup[pid] = [0, 0, -1, -1]
        l1, l2 = true_layer1[true_pid1 == pid], true_layer2[true_pid2 == pid]
        z1, z2 = true_z1[true_pid1 == pid], true_z2[true_pid2 == pid]
        for l in range(len(l1)):
            barrel_to_LEC = l1[l] in [0, 1, 2, 3] and l2[l] == 4
            barrel_to_REC = l1[l] in [0, 1, 2, 3] and l2[l] == 11
            if barrel_to_LEC or barrel_to_REC:
                temp = pid_lookup[pid]
                temp[0] += 1
                if abs(temp[1]) < abs(z1[l]):
                    if temp[0] > 1:

                        print(" *** adjusting y for triangle edge pattern")
                        print(
                            "     > replacing edge (l1={}, l2={}, z1={:.2f})\n".format(
                                temp[2], temp[3], temp[1]
                            ),
                            "      with      edge (l1={}, l2={}, z1={:.2f})".format(
                                l1[l], l2[l], abs(z1[l])
                            ),
                        )

                        y[
                            (pid1 == pid2)
                            & (pid1 == pid)
                            & (layer1 == temp[2])
                            & (layer2 == temp[3])
                        ] = 0

                    temp[1] = abs(z1[l])
                    temp[2] = l1[l]
                    temp[3] = l2[l]

                pid_lookup[pid] = temp

    print(" ... completed in {0} seconds".format(time.time() - t0))
    return Graph(X, edge_attr, edge_index, y, pid_mapped, pt)


def select_hits(hits, truth, particles, pt_min=0, endcaps=False):
    # Barrel volume and layer ids
    vlids = [(8, 2), (8, 4), (8, 6), (8, 8)]
    if endcaps:
        vlids.extend(
            [
                (7, 14),
                (7, 12),
                (7, 10),
                (7, 8),
                (7, 6),
                (7, 4),
                (7, 2),
                (9, 2),
                (9, 4),
                (9, 6),
                (9, 8),
                (9, 10),
                (9, 12),
                (9, 14),
            ]
        )
    n_det_layers = len(vlids)

    # Select barrel layers and assign convenient layer number [0-9]
    vlid_groups = hits.groupby(["volume_id", "layer_id"])
    hits = pd.concat(
        [vlid_groups.get_group(vlids[i]).assign(layer=i) for i in range(n_det_layers)]
    )
    # Calculate particle transverse momentum
    pt = np.sqrt(particles.px**2 + particles.py**2)
    particles["pt"] = pt
    # True particle selection.
    # Applies pt cut, removes all noise hits.
    particles = particles[pt > pt_min]
    truth = truth[["hit_id", "particle_id"]].merge(
        particles[["particle_id", "pt"]], on="particle_id"
    )
    # Calculate derived hits variables
    r = np.sqrt(hits.x**2 + hits.y**2)
    phi = np.arctan2(hits.y, hits.x)
    # Select the data columns we need
    hits = (
        hits[["hit_id", "z", "layer", "module_id"]]
        .assign(r=r, phi=phi)
        .merge(truth[["hit_id", "particle_id", "pt"]], on="hit_id")
    )
    # Remove duplicate hits (Skipped as per DeZoort's suggestion)
    # hits = hits.loc[hits.groupby(["particle_id", "layer"], as_index=False).r.idxmin()]

    return hits


def split_detector_sections(hits, phi_edges, eta_edges):
    """Split hits according to provided phi and eta boundaries."""
    hits_sections = []
    # Loop over sections
    for i in range(len(phi_edges) - 1):
        phi_min, phi_max = phi_edges[i], phi_edges[i + 1]
        # Select hits in this phi section
        phi_hits = hits[(hits.phi > phi_min) & (hits.phi < phi_max)]
        # Center these hits on phi=0
        centered_phi = phi_hits.phi - (phi_min + phi_max) / 2
        phi_hits = phi_hits.assign(phi=centered_phi, phi_section=i)
        for j in range(len(eta_edges) - 1):
            eta_min, eta_max = eta_edges[j], eta_edges[j + 1]
            # Select hits in this eta section
            eta = calc_eta(phi_hits.r, phi_hits.z)
            sec_hits = phi_hits[(eta > eta_min) & (eta < eta_max)]
            hits_sections.append(sec_hits.assign(eta_section=j))

    return hits_sections


def process_event(
    prefix,
    output_dir,
    module_maps,
    pt_min,
    n_eta_sections,
    n_phi_sections,
    eta_range,
    phi_range,
    phi_slope_max,
    z0_max,
    phi_reflect,
    endcaps,
    remove_intersecting_edges,
):
    # Load the data
    evtid = int(prefix[-9:])
    logging.info("Event %i, loading data" % evtid)
    hits, particles, truth = trackml.dataset.load_event(
        prefix, parts=["hits", "particles", "truth"]
    )

    # Apply hit selection
    logging.info("Event %i, selecting hits" % evtid)
    hits = select_hits(hits, truth, particles, pt_min=pt_min, endcaps=endcaps).assign(
        evtid=evtid
    )

    # Divide detector into sections
    # phi_range = (-np.pi, np.pi)
    phi_edges = np.linspace(*phi_range, num=n_phi_sections + 1)
    eta_edges = np.linspace(*eta_range, num=n_eta_sections + 1)
    hits_sections = split_detector_sections(hits, phi_edges, eta_edges)

    # Graph features and scale
    feature_names = ["r", "phi", "z"]
    feature_scale = np.array([1000.0, np.pi / n_phi_sections, 1000.0])
    if phi_reflect:
        feature_scale[1] *= -1

    # Define adjacent layers
    n_det_layers = 4
    l = np.arange(n_det_layers)
    layer_pairs = np.stack([l[:-1], l[1:]], axis=1)
    if endcaps:
        n_det_layers = 18
        EC_L = np.arange(4, 11)
        EC_L_pairs = np.stack([EC_L[:-1], EC_L[1:]], axis=1)
        layer_pairs = np.concatenate((layer_pairs, EC_L_pairs), axis=0)
        EC_R = np.arange(11, 18)
        EC_R_pairs = np.stack([EC_R[:-1], EC_R[1:]], axis=1)
        layer_pairs = np.concatenate((layer_pairs, EC_R_pairs), axis=0)
        barrel_EC_L_pairs = np.array([(0, 4), (1, 4), (2, 4)])  # , (3,4)])
        barrel_EC_R_pairs = np.array([(0, 11), (1, 11), (2, 11)])  # , (3,11)])
        layer_pairs = np.concatenate((layer_pairs, barrel_EC_L_pairs), axis=0)
        layer_pairs = np.concatenate((layer_pairs, barrel_EC_R_pairs), axis=0)

    # Construct the graph
    logging.info("Event %i, constructing graphs" % evtid)
    graphs = [
        construct_graph(
            section_hits,
            layer_pairs=layer_pairs,
            phi_slope_max=phi_slope_max,
            z0_max=z0_max,
            feature_names=feature_names,
            feature_scale=feature_scale,
            evtid=evtid,
            remove_intersecting_edges=remove_intersecting_edges,
            module_maps=module_maps,
        )
        for section_hits in hits_sections
    ]

    # Write these graphs to the output directory
    try:
        base_prefix = os.path.basename(prefix)
        filenames = [
            os.path.join(output_dir, "%s_g%03i" % (base_prefix, i))
            for i in range(len(graphs))
        ]
    except Exception as e:
        logging.info(e)

    logging.info("Event %i, writing graphs", evtid)
    for graph, filename in zip(graphs, filenames):
        np.savez(
            filename,
            **dict(
                x=graph.x,
                edge_attr=graph.edge_attr,
                edge_index=graph.edge_index,
                y=graph.y,
                pid=graph.pid,
            ),
        )


def main():
    """Main function"""

    # Parse the command line
    args = parse_args()

    # Setup logging
    log_format = "%(asctime)s %(levelname)s %(message)s"
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format=log_format)
    logging.info("Initializing")
    if args.show_config:
        logging.info("Command line config: %s" % args)

    # Load configuration
    with open(args.config) as f:
        config = yaml.load(f, yaml.FullLoader)
    if args.task == 0:
        logging.info("Configuration: %s" % config)

    # Load module maps
    pt_map = {
        0.5: "0p5",
        0.6: "0p6",
        0.7: "0p7",
        0.8: "0p8",
        0.9: "0p9",
        1: "1",
        1.1: "1p1",
        1.2: "1p2",
        1.3: "1p3",
        1.4: "1p4",
        1.5: "1p5",
        1.6: "1p6",
        1.7: "1p7",
        1.8: "1p8",
        1.9: "1p9",
        2.0: "2",
    }
    pt_str = pt_map[config["selection"]["pt_min"]]
    # indir_sample_str = [f for f in config["input_dir"].split("/") if "train" in f]
    # indir_sample_str = indir_sample_str[0].split("_")[1]
    map_sample_str = "000001000"  # "1" if indir_sample_str == "2" else "2"
    print(
        " *** using module maps from",
        f"module_maps/module_map_{map_sample_str}_{pt_str}GeV.npy",
    )
    module_maps = np.load(
        f"module_maps/module_map_{map_sample_str}_{pt_str}GeV.npy", allow_pickle=True
    ).item()
    module_maps = {key: item.astype(bool) for key, item in module_maps.items()}

    # Find the input files
    input_dir = config["input_dir"]
    all_files = os.listdir(input_dir)
    suffix = "-hits.csv"
    file_prefixes = sorted(
        os.path.join(input_dir, f.replace(suffix, ""))
        for f in all_files
        if f.endswith(suffix)
    )
    file_prefixes = file_prefixes[: config["n_files"]]

    # Split the input files by number of tasks and select my chunk only
    file_prefixes = np.array_split(file_prefixes, args.n_tasks)[args.task]
    file_prefixes = [
        prefix
        for prefix in file_prefixes
        if (
            (int(prefix.split("00000")[1]) >= args.start_evtid)
            and (int(prefix.split("00000")[1]) <= args.end_evtid)
        )
    ]

    # Prepare output
    output_dir = os.path.expandvars(config["output_dir"])
    os.makedirs(output_dir, exist_ok=True)
    logging.info("Writing outputs to " + output_dir)

    # Process input files with a worker pool
    t0 = time.time()
    with mp.Pool(processes=args.n_workers) as pool:
        process_func = partial(
            process_event,
            output_dir=output_dir,
            phi_range=(-np.pi, np.pi),
            module_maps=module_maps,
            **config["selection"],
        )
        pool.map(process_func, file_prefixes)

    t1 = time.time()
    print("Finished in", t1 - t0, "seconds")

    # Drop to IPython interactive shell
    if args.interactive:
        logging.info("Starting IPython interactive session")
        import IPython

        IPython.embed()

    logging.info("All done!")


if __name__ == "__main__":
    main()
