import os
import logging

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

from config import config
from utils import save_data
from visualize import visualize

# Takes too long. Tracks close to z-axis are scored high, other tracks are scored very low
# Track type 0: eps = 0.01
# Track type 1: eps = 0.008
def separate_track_types(df):
    for hit_id in df['hit_id']:
        hit = df[df['hit_id'] == hit_id]

        if hit.y.values > 0:
            if hit.z.values > 0:
                if hit.y.values < (1 / 3) * hit.z.values:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 0
                else:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 1
            else:
                if hit.y.values < - (1 / 3) * hit.z.values:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 0
                else:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 1
        else:
            if hit.z.values > 0:
                if hit.y.values > - (1 / 3) * hit.z.values:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 0
                else:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 1
            else:
                if hit.y.values > (1 / 3) * hit.z.values:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 0
                else:
                    df.loc[df['hit_id'].isin([hit_id]), 'track_type'] = 1

    return df

def cartesian_to_euclidean(hits):
    Rxy = np.sqrt(np.square(hits.x) + np.square(hits.y) +  np.square(hits.z))
    Rz = np.sqrt(np.square(hits.x) + np.square(hits.y))

    hits['x2'] = hits.x / Rxy
    hits['y2'] = hits.y / Rxy
    hits['z2'] = hits.z / Rz

    return hits

# from https://www.kaggle.com/c/trackml-particle-identification/discussion/57643#334429
def cartesian_to_cylindrical(df):
    df['rho'] = np.sqrt(np.square(df.x) + np.square(df.y))
    df['phi'] = np.arctan2(df.y, df.x)
    df['s'] = np.sin(df.phi)
    df['c'] = np.cos(df.phi)

    return df

# from https://www.kaggle.com/c/trackml-particle-identification/discussion/57643
def slice_cone(df, theta):
    df = df.assign(arctan2 = np.arctan2(df.z, df.rho))
    df = df.loc[(df.arctan2 > (theta-0.5)/180*np.pi) & (df.arctan2 < (theta+0.5)/180*np.pi)]

    return df

def view_positive_points(hits):
    hits = hits[hits.x > 0]
    hits = hits[hits.y > 0]
    hits = hits[hits.z > 0]

    return hits

def scale(df, z_scale):
    scaler = StandardScaler()
    scaler.fit(df[['x2', 'y2', 'z2']].values)

    df[['x2', 'y2', 'z2']] = scaler.transform(df[['x2', 'y2', 'z2']].values)

    df[['z']] = df[['z']] * z_scale

    return df

def select_volume(hits, particles, truth, volume):
    hits = hits[hits.volume_id == volume]
    truth = truth[truth.hit_id.isin(hits.hit_id.values)]
    particles = particles[particles.particle_id.isin(truth.particle_id.values)]

    return hits, particles, truth

def select_true_triples(hits, particles, truth):    
    triples = pd.DataFrame()

    hit_ids = hits.hit_id.values    
    for hit_id in hit_ids:
        particle_id = truth[truth.hit_id == hit_id]['particle_id'].values
        
        if particle_id != 0:
            hits_in_particle = truth[truth.particle_id.isin(particle_id)]['hit_id']

            if len(hits_in_particle) == 3:
                triple = hits[hits.hit_id.isin(hits_in_particle)][['hit_id','rho', 'phi']].values

                if triple.size == 0:
                    break

                triple_processed = pd.DataFrame()
                triple_processed['hit_id-1'] = [triple[0,0]] 
                triple_processed['hit_id-2'] = [triple[1,0]] 
                triple_processed['hit_id-3'] = [triple[2,0]] 
                triple_processed['rho-1'] = [triple[0,1]]
                triple_processed['phi-1'] = [triple[0,2]]
                triple_processed['rho-2'] = [triple[1,1]]
                triple_processed['phi-2'] = [triple[1,2]]
                triple_processed['rho-3'] = [triple[2,1]]
                triple_processed['phi-3'] = [triple[2,2]]
                
                triples = triples.append(triple_processed)
                hits = hits[-hits.hit_id.isin(hits_in_particle.values)]            

    return triples

def select_false_triples(hits, particles, truth):
    triples = pd.DataFrame()

    hit_ids = hits.hit_id.values    
    for hit_id in hit_ids:
        # find 3 points on different particles
        particle_id = truth[truth.hit_id == hit_id]['particle_id'].values
        
        other_particles = truth[truth.particle_id != particle_id[0]]
        other_particles = other_particles[other_particles.particle_id != 0]['particle_id']

        # if sample doesn't work get another sample
        other_particle_ids = other_particles.sample(2).values

        # get number of hits for each particle id
        hits_in_particle_1 = truth[truth.particle_id.isin(particle_id)]['hit_id']
        hits_in_particle_2 = truth[truth.particle_id == other_particle_ids[0]]['hit_id']            
        hits_in_particle_3 = truth[truth.particle_id == other_particle_ids[1]]['hit_id']

        # if the number of hits per particle is 3 continue
        if len(hits_in_particle_1) == 3 and len(hits_in_particle_2) == 3 and len(hits_in_particle_3) == 3:
            # get the hit_id for hit 1 in particle 1, hit 2 in particle 2 and hit 3 in particle 3
            point_1 = hits_in_particle_1.values[0]
            point_2 = hits_in_particle_2.values[1]
            point_3 = hits_in_particle_3.values[2]

            # change hit ids to rho and phi for each
            point_1 = hits[hits.hit_id == point_1][['hit_id','rho', 'phi']].values
            point_2 = hits[hits.hit_id == point_2][['hit_id','rho', 'phi']].values
            point_3 = hits[hits.hit_id == point_3][['hit_id','rho', 'phi']].values

            # package in df and append to triples
            triple_processed = pd.DataFrame()
            triple_processed['hit_id-1'] = [point_1[0,0]] 
            triple_processed['hit_id-2'] = [point_2[0,0]] 
            triple_processed['hit_id-3'] = [point_3[0,0]] 
            triple_processed['rho-1'] = [point_1[0,1]]
            triple_processed['phi-1'] = [point_1[0,2]]
            triple_processed['rho-2'] = [point_2[0,1]]
            triple_processed['phi-2'] = [point_2[0,2]]
            triple_processed['rho-3'] = [point_3[0,1]]
            triple_processed['phi-3'] = [point_3[0,2]]

            triples = triples.append(triple_processed)
            # hits_to_remove = triple_processed[['hit_id-1', 'hit_id-2', 'hit_id-3']].values[0].astype(int)
            # hits = hits[-hits.hit_id.isin(hits_to_remove)]
            # hit_ids = hits.hit_id.values    

    return triples

def combine_data(true_triples, false_triples):
    true_triples['true_track'] = 1
    false_triples['true_track'] = 0
    
    data = pd.DataFrame()

    data = data.append(true_triples)
    data = data.append(false_triples)

    return data

def preprocess(hits, cells=None, particles=None, truth=None, mode=None):
    logging.info('-' * 50)
    logging.info('Preprocessing')

    if config.EUCLIDEAN_TRANSFORM:
        logging.info('Transform hits to euclidean')
        hits = cartesian_to_euclidean(hits)

    if config.CYLINDRICAL_TRANSFORM:
        logging.info('Transform hits to cylindrical')
        hits = cartesian_to_cylindrical(hits)

    if config.CONE_SLICING:
        logging.info('Slicing cone at ' + str(config.THETA) + ' degrees')
        hits = slice_cone(hits, config.THETA)

    if config.VIEW_ONLY_POSITIVE_POINTS:
        logging.info('View positive points')
        hits = view_positive_points(hits)

    if config.SCALE:
        logging.info('Scaling')
        hits = scale(hits, config.Z_SCALE)

    if mode == 'TripleNet':
        logging.info('Selecting hits from volume')
        hits, particles, truth = select_volume(hits, particles, truth, 8)

        logging.info('Selecting groups of 3 hits on same track')
        true_triples = select_true_triples(hits, particles, truth)

        logging.info('Selecting groups of 3 hits on different tracks')
        false_triples = select_false_triples(hits, particles, truth)

        logging.info('Creating training data')
        data = combine_data(true_triples, false_triples)

        logging.info('-' * 50)
        logging.info('Saving data')
        save_data(hits, particles, truth, data, config.DATA_FOLDER)

        visualize(hits, particles, truth, data)
    else:
        return hits