# -*- coding: utf-8 -*-
"""
Created on Sun Feb 21 23:34:16 2016

Library of objects and internal functions

Ing.,Mgr. (MSc.) <NAME>
Biomedical engineering
International Clinical Research Center
St. Anne's University Hospital in Brno
Czech Republic
&
Mayo systems electrophysiology lab
Mayo Clinic
200 1st St SW
Rochester, MN
United States
"""

# %% Imports


import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# %% Node class


class CircNode:

    def __init__(self, position, label):
        self.position = position
        self.label = label
        self.group = None
        self.sublayer = None
        self.relative_value = 0
        self.color = "b"
        self.opacity = 1
        self.patch = None


# %% Connection class

class CircConn:

    def __init__(self, start_node, end_node):
        self.start_node = start_node
        self.end_node = end_node
        self.start_point = []
        self.end_point = []
        self.relative_value = 0
        self.color = "b"
        self.opacity = 1
        self.linewidth = 1
        self.patch = None

# %% Layer class


class CircLayer:

    def __init__(self, l_type, r, span, node_list=None, node_width=0,
                 conn_list=None):
        self.type = l_type
        self.r = r
        self.span = span
        self.node_list = node_list
        self.node_width = node_width
        self.conn_list = conn_list
        self.cmap = 'jet'  # Is this neccessary?
        self.sublayer = None  # This is only for heatmap
        self.labels = False


# %% Create labels

# TODO - make this a node funciton ?

def create_labels(node_list, node_width, ax):

    # Remove spines and grids
    ax.spines['polar'].set_visible(False)
    ax.grid(False)

    # Remove original ticks
    plt.xticks([])
    plt.yticks([])

    # Create new ticks
    tick_rotations = []
    tick_angles = []
    labels = []
    has = []
    for node in node_list:
        tick_angle = 180 * (node.position[0] + node_width/2) / np.pi
        tick_angles.append(tick_angle)
        labels.append(node.label)

        if tick_angle >= 270 or tick_angle <= 90:
            tick_rotations.append(tick_angle)
            has.append('left')
        else:
            # Flip the label, so text is always upright
            tick_rotations.append(tick_angle+180)
            has.append('right')

    ax.set_thetagrids(tick_angles, labels, frac=1.1)
    [x.set_rotation(tick_rotations[i]) for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_rotation_mode('anchor') for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_ha(has[i]) for i, x
     in enumerate(ax.get_xticklabels())]


def create_group_labels(node_list, node_width, ax):

    # Replicate axis
    ax = ax.figure.add_axes(ax.get_position(), projection='polar',
                            frameon=False)

    # Remove spines and grids
    ax.spines['polar'].set_visible(False)
    ax.grid(False)

    # Remove original ticks
    plt.xticks([])
    plt.yticks([])

    # Get groups
    group_list = list(set([node.group for node in node_list]))

    # Create new ticks
    tick_rotations = []
    tick_angles = []
    labels = []
    has = []
    for group in group_list:
        group_thetas = [node.position[0] for node in
                        node_list if node.group == group]
        group_start = min(group_thetas)
        group_end = max(group_thetas)
        group_width = group_end - group_start

        tick_angle = 180 * (group_start + group_width/2 + node_width/2) / np.pi
        tick_angles.append(tick_angle)
        labels.append(group)

        if tick_angle >= 270 or tick_angle <= 90:
            tick_rotations.append(tick_angle)
            has.append('left')
        else:
            # Flip the label, so text is always upright
            tick_rotations.append(tick_angle+180)
            has.append('right')

    max_nlabel_size = max([len(x.label) for x in node_list])
    ax.set_thetagrids(tick_angles, labels, frac=1.1+(max_nlabel_size*0.05))
    [x.set_rotation(tick_rotations[i]) for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_rotation_mode('anchor') for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_ha(has[i]) for i, x
     in enumerate(ax.get_xticklabels())]

# %% Node operations


def create_node_list(nodes, data, r, orderby=None, groupby=None,
                     sublayers=None, r_increment=0.1):

    # TODO - needs incremet for heatmap

    # Create groups if set
    N_groups = 0
    if groupby is not None:
        # Reorder the dataframe to match the groups
        if orderby is not None:
            data = data.sort_values(by=[groupby, orderby])
        else:
            data = data.sort_values(by=[groupby])

        group_list = list(data[groupby].unique())
        N_groups = len(group_list)  # This number is equal to number of gaps

    N = len(data[nodes].unique())  # This should be the number of channels
    thetas = np.linspace(0.0, 2*np.pi, N + N_groups, endpoint=False)
    node_width = 2 * np.pi / (N + N_groups)

    node_list = []
    for node_idx, node_label in enumerate(data[nodes].unique()):
        if groupby is not None:
            node_groups = data[(data[nodes] == node_label)][groupby].unique()
            if len(node_groups) > 1:
                print('Node belongs to more than one group!!!')
            group_idx = group_list.index(node_groups)

            node_idx += group_idx
            group = group_list[group_idx]
        else:
            group = None

        # FIXME: There has to be a better way to deal withe heatmap problem
        if sublayers is None:

            # Create node instance
            node = CircNode([thetas[node_idx], r], node_label)
            node.group = group

            # Add to node list
            node_list.append(node)
        else:
            i = 0
            for sublayer in data[sublayers].unique():
                row = r_increment * i
                # Create node instance
                node = CircNode([thetas[node_idx], r + row], node_label)
                node.group = group
                node.sublayer = sublayer

                i += 1

                # Add to node list
                node_list.append(node)

    return node_list, node_width


def reuse_nodes(node_list, nodes, data, r, sublayers=None,
                r_increment=0.1):

    # Take care of nodes coming from heatmap - takes only inner circle
    min_val = min([x.position[1] for x in node_list])
    node_list = [x for x in node_list if x.position[1] == min_val]

    new_node_list = []
    for node in node_list:

        # FIXME: There has to be a better way to deal withe heatmap problem
        if sublayers is None:

            # Create node instance
            new_node = CircNode([node.position[0], 0], node.label)
            new_node.group = node.group
            new_node.position[1] = r

            # Add to node list
            new_node_list.append(new_node)
        else:
            i = 0
            for sublayer in data[sublayers].unique():
                row = r_increment * i
                # Create node instance
                new_node = CircNode([node.position[0], 0], node.label)
                new_node.group = node.group
                new_node.sublayer = sublayer
                new_node.position[1] = r + row

                i += 1
                # Add to node list
                new_node_list.append(new_node)

    return new_node_list

# %% Connection operations


def create_conns(data, node_width, nodes, rel_interval):
    conns_df = pd.DataFrame()
    conns_df[['node_1', 'node_2', 'conn_val']] = data

    # Sort by connection strength? - to plot the strongest ones first?

    # Remove duplicate - move this into the function later
    conns_df = conns_df[(conns_df.node_1) != (conns_df.node_2)]
    max_value = conns_df.conn_val.max()
    min_value = conns_df.conn_val.min()
    val_range = max_value - min_value

    # Create relativistic values for connection strengths
    conns_df.conn_val = (conns_df.conn_val - min_value) / val_range

    # Reduce by interval
    if type(rel_interval[0]) == list:
        new_conn_df = pd.DataFrame()
        for interval in rel_interval:
            int_start = conns_df.conn_val.quantile(interval[0])
            int_stop = conns_df.conn_val.quantile(interval[1])
            pd.concat([new_conn_df,
                      conns_df[((conns_df.conn_val >= int_start) &
                                (conns_df.conn_val <= int_stop))]])
    else:
        int_start = conns_df.conn_val.quantile(rel_interval[0])
        int_stop = conns_df.conn_val.quantile(rel_interval[1])
        conns_df = conns_df[((conns_df.conn_val >= int_start) &
                             (conns_df.conn_val <= int_stop))]

    # Sorth the connections by strength
    conns_df = conns_df.sort_values(by=['conn_val'], ascending=True)

    # Get number of connections and reate start and end points with jitter
    noise_max = node_width * (len(nodes)*0.001)
    rng = np.random.mtrand.RandomState(seed=0)
    conns_df['node_1_conns'] = 1
    conns_df['node_2_conns'] = 1
    for node in nodes:
        conns_df.loc[(conns_df.node_1 == node.label),
                     'node_1_conns'] = sum(
                     conns_df.node_1 == node.label) + sum(
                     conns_df.node_2 == node.label)
        conns_df.loc[(conns_df.node_1 == node.label),
                     'node_2_conns'] = sum(
                     conns_df.node_2 == node.label) + sum(
                     conns_df.node_1 == node.label)

    for conn_row in conns_df.iterrows():
        i = conn_row[0]
        start_noise = rng.uniform(-noise_max, noise_max,
                                  conns_df.loc[i, 'node_1_conns'])[0]
        conns_df.loc[i, 'start_noise'] = start_noise
        end_noise = rng.uniform(-noise_max, noise_max,
                                conns_df.loc[i, 'node_2_conns'])[-1]
        conns_df.loc[i, 'end_noise'] = end_noise

    nodes_n_con_seen = np.zeros(len(nodes))
    conn_list = []
    for conn_row in conns_df.iterrows():
        i = conn_row[0]
        row_vals = conn_row[1]

        start_idxs = [idx for idx, x in enumerate(nodes) if x.label ==
                      row_vals.node_1]
        end_idxs = [idx for idx, x in enumerate(nodes) if x.label ==
                    row_vals.node_2]
        if len(start_idxs) == 0 or len(end_idxs) == 0:
            continue
        start = start_idxs[0]
        end = end_idxs[0]
        nodes_n_con_seen[start] += 1
        nodes_n_con_seen[end] += 1

        # Create connection object
        conn = CircConn([x for x in nodes if x.label == row_vals.node_1][0],
                        [x for x in nodes if x.label == row_vals.node_2][0])
        start_noise = conns_df.loc[i, 'start_noise']
        end_noise = conns_df.loc[i, 'end_noise']
        conn.start_point = ([conn.start_node.position[0] +
                            (node_width/2) +
                            start_noise,
                            conn.start_node.position[1]])
        conn.end_point = ([conn.end_node.position[0] +
                          (node_width/2) +
                          end_noise,
                          conn.end_node.position[1]])
        conn.relative_value = conns_df.loc[i, 'conn_val']
        conn_list.append(conn)

    return conn_list

# -*- coding: utf-8 -*-
"""
Created on Sun Feb 21 23:34:16 2016

Library of objects and internal functions

Ing.,Mgr. (MSc.) <NAME>
Biomedical engineering
International Clinical Research Center
St. Anne's University Hospital in Brno
Czech Republic
&
Mayo systems electrophysiology lab
Mayo Clinic
200 1st St SW
Rochester, MN
United States
"""

# %% Imports


import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# %% Node class


class CircNode:

    def __init__(self, position, label):
        self.position = position
        self.label = label
        self.group = None
        self.sublayer = None
        self.relative_value = 0
        self.color = "b"
        self.opacity = 1
        self.patch = None


# %% Connection class

class CircConn:

    def __init__(self, start_node, end_node):
        self.start_node = start_node
        self.end_node = end_node
        self.start_point = []
        self.end_point = []
        self.relative_value = 0
        self.color = "b"
        self.opacity = 1
        self.linewidth = 1
        self.patch = None

# %% Layer class


class CircLayer:

    def __init__(self, l_type, r, span, node_list=None, node_width=0,
                 conn_list=None):
        self.type = l_type
        self.r = r
        self.span = span
        self.node_list = node_list
        self.node_width = node_width
        self.conn_list = conn_list
        self.cmap = 'jet'  # Is this neccessary?
        self.sublayer = None  # This is only for heatmap
        self.labels = False


# %% Create labels

# TODO - make this a node funciton ?

def create_labels(node_list, node_width, ax):

    # Remove spines and grids
    ax.spines['polar'].set_visible(False)
    ax.grid(False)

    # Remove original ticks
    plt.xticks([])
    plt.yticks([])

    # Create new ticks
    tick_rotations = []
    tick_angles = []
    labels = []
    has = []
    for node in node_list:
        tick_angle = 180 * (node.position[0] + node_width/2) / np.pi
        tick_angles.append(tick_angle)
        labels.append(node.label)

        if tick_angle >= 270 or tick_angle <= 90:
            tick_rotations.append(tick_angle)
            has.append('left')
        else:
            # Flip the label, so text is always upright
            tick_rotations.append(tick_angle+180)
            has.append('right')

    ax.set_thetagrids(tick_angles, labels, frac=1.1)
    [x.set_rotation(tick_rotations[i]) for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_rotation_mode('anchor') for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_ha(has[i]) for i, x
     in enumerate(ax.get_xticklabels())]


def create_group_labels(node_list, node_width, ax):

    # Replicate axis
    ax = ax.figure.add_axes(ax.get_position(), projection='polar',
                            frameon=False)

    # Remove spines and grids
    ax.spines['polar'].set_visible(False)
    ax.grid(False)

    # Remove original ticks
    plt.xticks([])
    plt.yticks([])

    # Get groups
    group_list = list(set([node.group for node in node_list]))

    # Create new ticks
    tick_rotations = []
    tick_angles = []
    labels = []
    has = []
    for group in group_list:
        group_thetas = [node.position[0] for node in
                        node_list if node.group == group]
        group_start = min(group_thetas)
        group_end = max(group_thetas)
        group_width = group_end - group_start

        tick_angle = 180 * (group_start + group_width/2 + node_width/2) / np.pi
        tick_angles.append(tick_angle)
        labels.append(group)

        if tick_angle >= 270 or tick_angle <= 90:
            tick_rotations.append(tick_angle)
            has.append('left')
        else:
            # Flip the label, so text is always upright
            tick_rotations.append(tick_angle+180)
            has.append('right')

    max_nlabel_size = max([len(x.label) for x in node_list])
    ax.set_thetagrids(tick_angles, labels, frac=1.1+(max_nlabel_size*0.05))
    [x.set_rotation(tick_rotations[i]) for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_rotation_mode('anchor') for i, x
     in enumerate(ax.get_xticklabels())]
    [x.set_ha(has[i]) for i, x
     in enumerate(ax.get_xticklabels())]

# %% Node operations


def create_node_list(nodes, data, r, orderby=None, groupby=None,
                     sublayers=None, r_increment=0.1):

    # TODO - needs incremet for heatmap

    # Create groups if set
    N_groups = 0
    if groupby is not None:
        # Reorder the dataframe to match the groups
        if orderby is not None:
            data = data.sort_values(by=[groupby, orderby])
        else:
            data = data.sort_values(by=[groupby])

        group_list = list(data[groupby].unique())
        N_groups = len(group_list)  # This number is equal to number of gaps

    N = len(data[nodes].unique())  # This should be the number of channels
    thetas = np.linspace(0.0, 2*np.pi, N + N_groups, endpoint=False)
    node_width = 2 * np.pi / (N + N_groups)

    node_list = []
    for node_idx, node_label in enumerate(data[nodes].unique()):
        if groupby is not None:
            node_groups = data[(data[nodes] == node_label)][groupby].unique()
            if len(node_groups) > 1:
                print('Node belongs to more than one group!!!')
            group_idx = group_list.index(node_groups)

            node_idx += group_idx
            group = group_list[group_idx]
        else:
            group = None

        # FIXME: There has to be a better way to deal withe heatmap problem
        if sublayers is None:

            # Create node instance
            node = CircNode([thetas[node_idx], r], node_label)
            node.group = group

            # Add to node list
            node_list.append(node)
        else:
            i = 0
            for sublayer in data[sublayers].unique():
                row = r_increment * i
                # Create node instance
                node = CircNode([thetas[node_idx], r + row], node_label)
                node.group = group
                node.sublayer = sublayer

                i += 1

                # Add to node list
                node_list.append(node)

    return node_list, node_width


def reuse_nodes(node_list, nodes, data, r, sublayers=None,
                r_increment=0.1):

    # Take care of nodes coming from heatmap - takes only inner circle
    min_val = min([x.position[1] for x in node_list])
    node_list = [x for x in node_list if x.position[1] == min_val]

    new_node_list = []
    for node in node_list:

        # FIXME: There has to be a better way to deal withe heatmap problem
        if sublayers is None:

            # Create node instance
            new_node = CircNode([node.position[0], 0], node.label)
            new_node.group = node.group
            new_node.position[1] = r

            # Add to node list
            new_node_list.append(new_node)
        else:
            i = 0
            for sublayer in data[sublayers].unique():
                row = r_increment * i
                # Create node instance
                new_node = CircNode([node.position[0], 0], node.label)
                new_node.group = node.group
                new_node.sublayer = sublayer
                new_node.position[1] = r + row

                i += 1
                # Add to node list
                new_node_list.append(new_node)

    return new_node_list

# %% Connection operations


def create_conns(data, node_width, nodes, rel_interval):
    conns_df = pd.DataFrame()
    conns_df[['node_1', 'node_2', 'conn_val']] = data

    # Sort by connection strength? - to plot the strongest ones first?

    # Remove duplicate - move this into the function later
    conns_df = conns_df[(conns_df.node_1) != (conns_df.node_2)]
    max_value = conns_df.conn_val.max()
    min_value = conns_df.conn_val.min()
    val_range = max_value - min_value

    # Create relativistic values for connection strengths
    conns_df.conn_val = (conns_df.conn_val - min_value) / val_range

    # Reduce by interval
    if type(rel_interval[0]) == list:
        new_conn_df = pd.DataFrame()
        for interval in rel_interval:
            int_start = conns_df.conn_val.quantile(interval[0])
            int_stop = conns_df.conn_val.quantile(interval[1])
            pd.concat([new_conn_df,
                      conns_df[((conns_df.conn_val >= int_start) &
                                (conns_df.conn_val <= int_stop))]])
    else:
        int_start = conns_df.conn_val.quantile(rel_interval[0])
        int_stop = conns_df.conn_val.quantile(rel_interval[1])
        conns_df = conns_df[((conns_df.conn_val >= int_start) &
                             (conns_df.conn_val <= int_stop))]

    # Sorth the connections by strength
    conns_df = conns_df.sort_values(by=['conn_val'], ascending=True)

    # Get number of connections and reate start and end points with jitter
    noise_max = node_width * (len(nodes)*0.001)
    rng = np.random.mtrand.RandomState(seed=0)
    conns_df['node_1_conns'] = 1
    conns_df['node_2_conns'] = 1
    for node in nodes:
        conns_df.loc[(conns_df.node_1 == node.label),
                     'node_1_conns'] = sum(
                     conns_df.node_1 == node.label) + sum(
                     conns_df.node_2 == node.label)
        conns_df.loc[(conns_df.node_1 == node.label),
                     'node_2_conns'] = sum(
                     conns_df.node_2 == node.label) + sum(
                     conns_df.node_1 == node.label)

    for conn_row in conns_df.iterrows():
        i = conn_row[0]
        start_noise = rng.uniform(-noise_max, noise_max,
                                  conns_df.loc[i, 'node_1_conns'])[0]
        conns_df.loc[i, 'start_noise'] = start_noise
        end_noise = rng.uniform(-noise_max, noise_max,
                                conns_df.loc[i, 'node_2_conns'])[-1]
        conns_df.loc[i, 'end_noise'] = end_noise

    nodes_n_con_seen = np.zeros(len(nodes))
    conn_list = []
    for conn_row in conns_df.iterrows():
        i = conn_row[0]
        row_vals = conn_row[1]

        start_idxs = [idx for idx, x in enumerate(nodes) if x.label ==
                      row_vals.node_1]
        end_idxs = [idx for idx, x in enumerate(nodes) if x.label ==
                    row_vals.node_2]
        if len(start_idxs) == 0 or len(end_idxs) == 0:
            continue
        start = start_idxs[0]
        end = end_idxs[0]
        nodes_n_con_seen[start] += 1
        nodes_n_con_seen[end] += 1

        # Create connection object
        conn = CircConn([x for x in nodes if x.label == row_vals.node_1][0],
                        [x for x in nodes if x.label == row_vals.node_2][0])
        start_noise = conns_df.loc[i, 'start_noise']
        end_noise = conns_df.loc[i, 'end_noise']
        conn.start_point = ([conn.start_node.position[0] +
                            (node_width/2) +
                            start_noise,
                            conn.start_node.position[1]])
        conn.end_point = ([conn.end_node.position[0] +
                          (node_width/2) +
                          end_noise,
                          conn.end_node.position[1]])
        conn.relative_value = conns_df.loc[i, 'conn_val']
        conn_list.append(conn)

    return conn_list

# -*- coding: utf-8 -*- Created on Sun Feb 21 23:34:16 2016  Library of objects and internal functions  Ing.,Mgr. (MSc.) <NAME> Biomedical engineering International Clinical Research Center St. Anne's University Hospital in Brno Czech Republic & Mayo systems electrophysiology lab Mayo Clinic 200 1st St SW Rochester, MN United States # %% Imports # %% Node class # %% Connection class # %% Layer class # Is this neccessary? # This is only for heatmap # %% Create labels # TODO - make this a node funciton ? # Remove spines and grids # Remove original ticks # Create new ticks # Flip the label, so text is always upright # Replicate axis # Remove spines and grids # Remove original ticks # Get groups # Create new ticks # Flip the label, so text is always upright # %% Node operations # TODO - needs incremet for heatmap # Create groups if set # Reorder the dataframe to match the groups # This number is equal to number of gaps # This should be the number of channels # FIXME: There has to be a better way to deal withe heatmap problem # Create node instance # Add to node list # Create node instance # Add to node list # Take care of nodes coming from heatmap - takes only inner circle # FIXME: There has to be a better way to deal withe heatmap problem # Create node instance # Add to node list # Create node instance # Add to node list # %% Connection operations # Sort by connection strength? - to plot the strongest ones first? # Remove duplicate - move this into the function later # Create relativistic values for connection strengths # Reduce by interval # Sorth the connections by strength # Get number of connections and reate start and end points with jitter # Create connection object