#!/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)


def getXRange(params):
    O1P = params['acqus']['O1']/params['acqus']['SFO1']
    SW = params['acqus']['SW']

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

    return O1P + SW / 2, O1P - SW / 2


def getYRange(pdata, x_ticks):
    y_sum = pdata.sum(axis=1)

    # calculcate deltaC based on known mobility
    #c = calcMobility(angle=0.6875, distance=0.0377, bigDelta=(300 * 10 ** -3), gradPower=35 * 10 ** -3, smallDelta=(10 ** -3)) / 10
    #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
    return (0 - zero_index) / deltaC, (pdata.shape[0] - 1 - zero_index) / deltaC


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

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'))

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

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

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

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

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

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()