#!/usr/bin/env python3

import matplotlib.pyplot as plt
import nmrglue
import numpy as np
import os
import sys


def getMaxIndex(pdata, x_ticks, position, epsilon=1.0):
    indexes = np.where((x_ticks > position - epsilon) & (x_ticks < position + epsilon))[0]
    return np.argmax(pdata[:,indexes].sum(axis=1))


def calcMobility(angle, distance, bigDelta, gradPower, smallDelta):
    # bigDelta - d20
    # smallDelta - p30
    return angle * distance * np.pi * 10 ** 8 / (bigDelta * (2.68 * 10 ** 8) * gradPower * 180 * smallDelta)


left, width = 0.1, 0.65
bottom, height = 0.1, 0.65
spacing = 0.005
spacing = 0.03

c = calcMobility(angle=0.6875, distance=0.0377, bigDelta=(300 * 10 ** -3), gradPower=35 * 10 ** -3, smallDelta=(10 ** -3)) / 10

countour = [left, bottom, width, height]
spectrum_x = [left, bottom + height + spacing, width, 0.2]
#spectrum_y = [left + width + spacing, bottom, 0.2, height]

plt.figure(figsize = (15,8))

countour = plt.axes(countour)
countour.tick_params(direction='in', top=True, right=True)
spectrum_x = plt.axes(spectrum_x)
spectrum_x.tick_params(direction='in', labelbottom=False)
#spectrum_y = plt.axes(spectrum_y)
#spectrum_y.tick_params(direction='in', labelleft=False)


path = os.path.abspath(sys.argv[1] if len(sys.argv) == 2 else os.getcwd())
params, pdata = nmrglue.fileio.bruker.read_pdata(dir=os.path.join(path, 'pdata', '1'))

print(nmrglue.fileio.bruker.guess_udic(params, pdata))

#for k,v in params.items():
#    print(str(k) + ':' + str(v))
#    print('')

shape = pdata.shape
O1P = params['acqus']['O1']/params['acqus']['SFO1']
SW = params['acqus']['SW']

Y_O1P = params['acqu2s']['O1']/params['acqu2s']['SFO1']
Y_SW = params['acqu2s']['SW']

# X, Y = np.meshgrid(np.linspace(O1P+SW/2, O1P-SW/2, shape[1]), np.arange(shape[0]))
x_min, x_max = O1P + SW / 2, O1P - SW / 2

x_sum = pdata.sum(axis=0)
y_sum = pdata.sum(axis=1)

x_ticks = np.linspace(x_min, x_max, shape[1])

# calculcate deltaC based on known mobility
max_index = np.argmax(y_sum)
zero_index = getMaxIndex(pdata, x_ticks, 4.79)
deltaC = (max_index - zero_index) / c
print(deltaC)

# use hardcoded deltaC when mobility is not known
#zero_index = np.argmax(y_sum)
#deltaC = 2.4882422732549014
y_min, y_max = (0 - zero_index) / deltaC, (pdata.shape[0] - 1 - zero_index) / deltaC

y_ticks = np.linspace(y_min, y_max, shape[0])

X, Y = np.meshgrid(x_ticks, y_ticks)

levels = np.linspace(np.min(pdata), np.max(pdata), 50)

countour.contourf(X, Y, pdata, levels=levels[2:5])
#countour.contour(X, Y, pdata, levels=20)
countour.set_xlabel('Chemical shift, ppm')
countour.set_xlim(x_min, x_max)
countour.set_ylim(y_min, y_max)
countour.set_ylabel('Electrophoretic Mobility (10 ** -8)')

spectrum_x.plot(x_ticks, x_sum)
spectrum_x.set_xlim(x_min, x_max)
spectrum_x.axis('off')

#spectrum_y.plot(y_sum, y_ticks)
#spectrum_y.set_ylim(y_min, y_max)
#spectrum_y.axis('off')

plt.show()