#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 27 13:05:38 2019

@author: ltlustos
"""

import numpy as np
import pylab as pl
# import sympy
import matplotlib.pyplot as plt
import matplotlib.ticker

# (x,y, x0,y0, u, X, Y, Q, Q0, s, enc, pitch, pi, sqrt2) = sympy.symbols('x y x0 y0 u X Y Q Q0 s enc pitch pi sqrt2')
# s_gauss_x = 1.0 / s / sqrt2 /sympy.sqrt(pi) * sympy.exp(-0.5 * (x-x0)**2 / s **2)
# s_gauss_y = 1.0 / s / sqrt2 / sympy.sqrt(pi) * sympy.exp(-0.5 * (y-y0)**2 / s **2)
# f1 = sympy.integrate(sympy.integrate(s_gauss_x*s_gauss_y , (x, -pitch/2, pitch/2)), (y, -pitch/2, pitch/2))

datapath = '/Users/marialauraperez/Desktop/SummerCERN/Results/'
dataout = '/Users/marialauraperez/Desktop/SummerCERN/Results/Charge collection/'

def pixel_charge(x0, y0, s, pitch):
    from scipy.special import erf
    from numpy import pi
    _sqrt2 = np.sqrt(2.0)
    return 0.5*pi*erf(-(1/_sqrt2)*(-pitch/2 - x0)/s)*erf(-(1/_sqrt2)*(-pitch/2 - y0)/s)/(pi*_sqrt2**2) + \
           0.5*pi*erf(-(1/_sqrt2)*(-pitch/2 - x0)/s)*erf((1/_sqrt2)*(pitch/2 - y0)/s)/(pi*_sqrt2**2) + \
           0.5*pi*erf(-(1/_sqrt2)*(-pitch/2 - y0)/s)*erf((1/_sqrt2)*(pitch/2 - x0)/s)/(pi*_sqrt2**2) + \
           0.5*pi*erf((1/_sqrt2)*(pitch/2 - x0)/s)*erf((1/_sqrt2)*(pitch/2 - y0)/s)/(pi*_sqrt2**2)


def normalize(arr):
    norm = np.linalg.norm(arr)
    return arr/norm


def rex_el(kv, Z, A, rhocm3):
    # ex_el(kv, Z, A, varargin)
    # um
    # NUCLEAR INSTRUMENTS AND METHODS IO3 (I972) 85-91;
    # GENERALIZED SEMIEMPIRICAL EQUATIONS
    # FOR THE EXTRAPOLATED RANGE OF ELECTRONS
    # T. TABATA, R. ITO and S. OKABE
    b=[.2335, 1.209, 1.78e-4, .9891, 3.01e-4, 1.468, 1.180e-2, 1.232,.109]
    a= [b[0]*A/Z**b[1], b[2]*Z, b[3]-b[4]*Z, b[5]-b[6]*Z, b[7]/Z**b[8]]
    tau= kv/511.03
    rex=a[0]*(np.log(1+a[1]*tau)/a[1]-a[2]*tau/(1+a[3]*tau**a[4]))/rhocm3*1e4 # um
    return rex


pitch_um = 55.0
s_um = np.array([0.71358728, 1.01030249, 1.238763, 1.43202509, 1.60287829, 1.75787494, 1.90090355, 2.03449532,
                 2.16040679, 2.27991584, 2.39398634, 2.50336657, 2.60865116, 2.71032194, 2.80877575, 2.90434405,
                 2.99730709, 3.08790435, 3.17634237, 3.26280079, 3.34743701, 3.43038988, 3.51178263, 3.59172523,
                 3.67031633, 3.7476448, 3.8237911, 3.89882831, 3.97282309, 4.04583643, 4.11792434, 4.18913838,
                 4.25952615, 4.32913174, 4.39799607, 4.46615721, 4.53365067, 4.60050961, 4.66676511, 4.73244631,
                 4.79758063, 4.86219387, 4.92631037, 4.98995314, 5.05314394, 5.11590341, 5.17825114, 5.24020575,
                 5.30178498, 5.36300571, 5.42388409, 5.48443553, 5.54467478, 5.60461599, 5.66427272, 5.72365798,
                 5.78278432, 5.84166378, 5.90030798, 5.95872814, 6.01693507, 6.07493925, 6.13275081, 6.19037957,
                 6.24783503, 6.30512645, 6.3622628, 6.41925282, 6.47610502, 6.53282769, 6.58942892, 6.64591661,
                 6.70229848, 6.75858209, 6.81477483, 6.87088397, 6.92691662, 6.98287977, 7.03878029, 7.09462495,
                 7.1504204, 7.20617322, 7.26188987, 7.31757675, 7.37324019, 7.42888644, 7.48452169, 7.54015208,
                 7.59578371, 7.65142261, 7.70707479, 7.76274622, 7.81844286, 7.87417062, 7.92993541, 7.98574313,
                 8.04159967, 8.0975109, 8.15348272])  # taken from the output of Sensor.py

energies = np.linspace(4, 60, 113)  # spacing 0.5
selected_energies = np.linspace(5, 50, 10)
selected_weights = list()
all_weights = np.genfromtxt(datapath+'sigma_weights.txt')
for ii in range(len(energies)):
    if energies[ii] in selected_energies:
        selected_weights.append(all_weights[:, ii])


bins = np.linspace(0, 1, 100)
n_vals = []
for e in range(len(selected_energies)):
    n_vals.append([])
    for i in range(len(s_um)):
        X, Y = np.meshgrid(np.linspace(-pitch_um/2.0, pitch_um/2.0, 500), np.linspace(-pitch_um/2.0, pitch_um/2.0, 500))
        Q = pixel_charge(X, Y, s_um[i] + rex_el(e, 14, 28.08, 2.32), pitch_um)
        n, bins, patches = pl.hist(Q.ravel(), bins=bins, histtype='step')
        n_vals[e].append(n)
bins2 = (bins[1:] + bins[:-1]) / 2.0
f, ax = plt.subplots(1, 1, figsize=(8, 6))
for e in range(len(selected_energies)):
    total_spectra = np.zeros(len(n_vals[e][0]))
    for k in range(len(n_vals[e][0])):
        for j in range(len(s_um)):
            total_spectra[k] += n_vals[e][j][k]*selected_weights[e][j]
    plt.plot(bins2, normalize(total_spectra)*100, label=r'$E_\gamma$= '+str(selected_energies[e])+' keV')
plt.xlabel('Charge collection efficiency')
plt.ylabel('Percentage of occurrence')
plt.yscale('log')
plt.legend(fontsize='x-small')
plt.xlim(left=0.2)
ax.get_yaxis().set_major_formatter(matplotlib.ticker.PercentFormatter(decimals=2))
plt.title('Total charge collection efficiency within a pixel', fontweight='bold')
plt.savefig(dataout+'Charge_efficiency2.png')