# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + _cell_guid="79c7e3d0-c299-4dcb-8224-4455121ee9b0" _uuid="d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
# %matplotlib inline

import os
import multiprocessing
from multiprocessing import Pool

import numpy as np
import pandas as pd
from tqdm import tqdm
import mlcrate as mlc
from bayes_opt import BayesianOptimization

from trackml.dataset import load_event
from trackml.dataset import load_dataset
from trackml.score import score_event

from sklearn.utils import shuffle
from sklearn.preprocessing import StandardScaler
from sklearn.cluster.dbscan_ import dbscan

import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# import hdbscan

class Config():
    ROOT_PATH = '/home/bilal/.kaggle/competitions/trackml-particle-identification'
    
    DATA_FOLDERS = {
        'train': 'train_100_events', 
        'test': 'test'
    }
    
    EVENT = 'event000001000'

    OUTPUT_FOLDER = './data'
    
    SPLIT = 0.75
    THETA = 75

config = Config()

event_path = os.path.join(config.ROOT_PATH, config.DATA_FOLDERS['train'], config.EVENT)
hits, cells, particles, truth = load_event(event_path)


# + _uuid="a6033df117e0b8800f9572208c37ffa5ccf653ff"
def preprocess(hits):
    hits['r1'] = np.sqrt(np.square(hits['x']) + np.square(hits['y']))
    hits['r2'] = np.sqrt(np.square(hits['x']) + np.square(hits['y']) + np.square(hits['z']))
    hits['x2'] = hits['x'] / hits['r1']
    hits['y2'] = hits['y'] / hits['r1']
    hits['z1'] = hits['z'] / hits['r1']
    hits['z2'] = hits['z'] / hits['r2']

    hits['rho'] = np.sqrt(np.square(hits['x']) + np.square(hits['y']))
    hits['phi'] = np.arctan2(hits['y'], hits['x'])
    hits['s'] = np.sin(hits['phi'])
    hits['c'] = np.cos(hits['phi'])
    
    hits['arctan2'] = np.arctan2(hits['z'], hits['rho'])
    
    hits['log1p'] = np.log1p(np.abs(hits['z2'])) * np.sign(hits['z2'])
    
    hits['r-abs'] = (hits['r2'] - np.abs(hits['z'])) / hits['r2']
    
    return hits

def get_preds(x):
    w1, w2, w3, param = x[0], x[1], x[2], x[3]
    hits, mm, ii = param[0], param[1], param[2]

    hits['aa'] = hits['phi'] + mm * (hits['r1'] + 0.000005 * np.square(hits['r1'])) / 1000 * (ii / 2) / 180 * np.pi
    hits['cosaa'] = np.cos(hits['aa'])
    hits['sinaa'] = np.sin(hits['aa'])
    
    x = StandardScaler().fit_transform(np.column_stack([hits['cosaa'], hits['sinaa'], hits['z1'], hits['log1p']]))

    x[:,0] = x[:,0] * w1
    x[:,1] = x[:,1] * w1
    x[:,2] = x[:,2] * w2
    x[:,3] = x[:,3] * w3

    # Change to 4 if needed
    _, preds = dbscan(x, eps=0.0035, min_samples=1, n_jobs=8)

    return preds

def get_params(hits, niter):
    params = []
    for i in range(0, int(niter)):
        ii = i

        if i % 2 == 0:
            mm = 1
        else:
            mm = -1
            
        params.append((hits, mm, ii))
        
    return params

def train(hits, w1, w2, w3, niter, optimize=False):
    params = get_params(hits, niter)
    
    for i, param in enumerate(params):
        params[i] = (w1, w2, w3, param)
    
    # Change to 4 if needed
    pool = Pool(processes=8)
    preds = pool.map(get_preds, params)
    pool.close()

    return preds

def add_count(l):
    # unique: sorted unique values; reverse: indicies of unique to reconstruct l; count: num times each unique appears
    unique, reverse, count = np.unique(l, return_counts=True, return_inverse=True)
        
    # get num times each unique l appears
    c = count[reverse]
    
    # unassign any tracks with either 0 or > 20 hits    
    c[np.where(l == 0)] = 0
    c[np.where(c > 20)] = 0
    
    return (l, c)

def postprocess(hits, preds, event):
    results = [add_count(l) for l in preds]

    preds, counts = results[0]

    for i in range(1, len(results)):
        l, c = results[i]
        idx = np.where((c - counts > 0))[0]
        preds[idx] = l[idx] + preds.max()
        counts[idx] = c[idx]

    hits['track_id'] = preds
    hits['event_id'] = event
    
    return hits

def predict(w1, w2, w3, niter=10, optimize=True, visualize=False, test=False, hits=hits, event='000001000'):
    hits = preprocess(hits)
    preds = train(hits, w1, w2, w3, niter, optimize=optimize)
    hits = postprocess(hits, preds, event)
    
    if not test:
        score = score_event(truth, hits)
    
    if visualize:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        hits = hits[hits['track_id'] != 0]

        for particle in hits['track_id'][:100].unique():
            hit = hits[hits['track_id'] == particle]

            ax.scatter(hit.x, hit.y, hit.z, marker='o')
            ax.plot(hit.x, hit.y, hit.z)

        ax.set_xlabel('X Label')
        ax.set_ylabel('Y Label')
        ax.set_zlabel('Z Label')

        plt.show()
    
    if test:
        return hits
    else:
        if optimize:
            return score
        else:
            return hits, score


# -

hits, score = predict(1, 1, 1, niter=100, optimize=False, visualize=True)
print(score)

hits, score = predict(0.5, 1, 1, niter=100, optimize=False, visualize=True)
print(score)

hits, score = predict(1, 0.5, 1, niter=100, optimize=False, visualize=True)
print(score)

hits, score = predict(0.56, 0.41, 0.6, niter=167, optimize=False, visualize=True)
print(score)

hits, score = predict(0.56, 0.41, 0.6, niter=250, optimize=False, visualize=True)
print(score)

# Step |   Time |      Value |     niter |        w1 |        w2 | w3 |
# ------|--------|------------|-----------|-----------|-----------|---|
#     5 | 00m46s |    0.45445 |  167.3082 |    0.5689 |    0.4154 |    0.5990 | 
#     2 | 00m39s |    0.44926 |  137.5896 |    0.5283 |    0.6565 |    0.5012 | 
#     1 | 00m59s |    0.45261 |  203.5614 |    0.5499 |    0.4267 |    0.5110 | 
#     2 | 00m58s |    0.45268 |  216.4917 |    0.7397 |    0.6608 |    0.7940 | 
#     5 | 00m53s |    0.46916 |  194.8645 |    0.7772 |    0.5834 |    0.5569 | 
#     8 | 00m55s |    0.46961 |  199.6406 |    1.0376 |    0.6162 |    0.4525 | 
#     5 | 01m02s |    0.48692 |  235.2411 |    1.0453 |    0.4246 |    0.4109 | 
#    12 | 01m16s |    0.48994 |  250.0000 |    1.2000 |    0.4000 |    0.3000 | 
#     6 | 01m12s |    0.49100 |  266.7720 |    1.3420 |    0.4187 |    0.2252 | 
#     4 | 01m13s |    0.49135 |  272.3951 |    1.0503 |    0.4926 |    0.2475 | 
#     1 | 01m13s |    0.49147 |  262.6490 |    1.0993 |    0.4470 |    0.2170 | 
#     2 | 01m00s |    0.49199 |  204.0260 |    1.0781 |    0.4426 |    0.2325 | 
#

bo = BayesianOptimization(predict, {'w1': (1, 1.2), 'w2': (0.4, 0.5), 'w3': (0.15, 0.25), 'niter': (200, 300)})

bo.initialize(
    {
        'target': [0.49147, 0.49135],
        'niter': [262, 272],
        'w1': [1.0993, 1.0503],
        'w2': [0.4470, 0.4926],
        'w3': [0.2170, 0.2475]
    }
)

bo.maximize(init_points=5, n_iter=10, kappa=2)

# +
submissions = []

for event_id, hits, cells in tqdm(load_dataset(os.path.join(config.ROOT_PATH, config.DATA_FOLDERS['test']), parts=['hits', 'cells'])):
# for event_id, hits, cells in tqdm(load_dataset(os.path.join(config.ROOT_PATH, config.DATA_FOLDERS['test']), parts=['hits', 'cells'], nevents=2)):
    hits = predict(1.0438, 0.3795, 0.2350, niter=214, optimize=False, visualize=False, test=True, hits=hits, event=event_id)
    
    sub = hits[['hit_id', 'track_id', 'event_id']]
 
    submissions.append(sub)
# -

submission = pd.concat(submissions, axis=0)

len(submission)

mlc.kaggle.save_sub(submission, 'submission.csv.gz')

submission

StarcoderdataJupyterScriptsDedupFiltered