import numpy as np
import matplotlib.pyplot as plt

import jetFunctions as jet

dataPath = '11-jet/data/'
plotsPath = '11-jet/plots/'


########################################
#   Pressure calibration
########################################

p00, steps, count = jet.read(dataPath + 'calibration/00 Pa.txt')
p70, steps, count = jet.read(dataPath + 'calibration/70 Pa.txt')

pressure = [0, 70] # Pa
adc = [int(p00.mean()), int(p70.mean())]

linearFit = np.polyfit(adc, pressure, 1)

adcFit = np.linspace(int(p00.mean()), int(p70.mean()), 100)
pressureFit = np.polyval(linearFit, adcFit)


# Plot

plt.figure('Калибровка давления', figsize=(8, 6), dpi=200)

plt.title('Калибровочный график\nзависимости показаний АЦП от давления')
plt.xlabel('Далвение [Па]')
plt.ylabel('Отсчёты АЦП')

plt.plot(pressure, adc, '*', label='Измерения')
plt.plot(pressureFit, adcFit, label='P = {:.3f} * N - {:.2f} [Па]'.format(linearFit[0], -linearFit[1]))
plt.legend()

plt.grid(which='major')
plt.grid(which='minor', linestyle='--')
plt.minorticks_on()


plt.savefig(plotsPath + 'pressure-calibration.png')


########################################
#   Distance calibration
########################################

xTotalDistance = 50 / 1000 # m
xTotalSteps = 900 # steps
xStep = xTotalDistance / xTotalSteps # m / step

step = [0, xTotalSteps]
distance = np.array([0, xTotalDistance])

stepsLinearFit = np.polyfit(step, distance, 1)

stepFit = np.linspace(0, xTotalSteps, 100)
distanceFit = np.polyval(stepsLinearFit, stepFit)

# Plot

plt.figure('Калибровка расстояния', figsize=(8, 6), dpi=200)

plt.title('Калибровочный график\nзависимости перемещения трубки Пито от шага двигателя')
plt.xlabel('Количество шагов')
plt.ylabel('Перемещение трубки Пито [см]')

plt.plot(step, distance * 100, '*', label='Измерения')
plt.plot(stepFit, distanceFit * 100, label='X = {:.1e} * step [м]'.format(xStep))
plt.legend()

plt.grid(which='major')
plt.grid(which='minor', linestyle='--')
plt.minorticks_on()

plt.savefig(plotsPath + 'distance-calibration.png')


########################################
#   Processing
########################################

measuresDirectory = '100 measures/'

airDensity = 1.2 # kg / m^3

def calculateVelocity(filename, xCorrection):
    data, steps, count = jet.read(filename)

    p = np.polyval(linearFit, data) # Pa
    p[p < 3] = 0
    u = np.sqrt(2 * abs(p) / airDensity) # m / s
    
    dx = xStep * steps
    x = np.linspace(0, count * dx, count) - xCorrection / 1000

    return x, u, dx

def calculateOutgo(x, u, dx):
    dq = 2 * np.pi * airDensity * u * abs(x) * dx # kg / s
    q = sum(dq) / 2 # kg / s
    
    return q

x00, u00, dx = calculateVelocity(dataPath + measuresDirectory + '00 mm.txt', 27.9)
x10, u10, dx = calculateVelocity(dataPath + measuresDirectory + '10 mm.txt', 27.9)
x20, u20, dx = calculateVelocity(dataPath + measuresDirectory + '20 mm.txt', 27.9)
x30, u30, dx = calculateVelocity(dataPath + measuresDirectory + '30 mm.txt', 27.4)
x40, u40, dx = calculateVelocity(dataPath + measuresDirectory + '40 mm.txt', 28.5)
x50, u50, dx = calculateVelocity(dataPath + measuresDirectory + '50 mm.txt', 26.9)
x60, u60, dx = calculateVelocity(dataPath + measuresDirectory + '60 mm.txt', 26.5)
x70, u70, dx = calculateVelocity(dataPath + measuresDirectory + '70 mm.txt', 27.5)
x80, u80, dx = calculateVelocity(dataPath + measuresDirectory + '80 mm.txt', 28.9)
x90, u90, dx = calculateVelocity(dataPath + measuresDirectory + '90 mm.txt', 31.2)

q00 = calculateOutgo(x00, u00, dx)
q10 = calculateOutgo(x10, u10, dx)
q20 = calculateOutgo(x20, u20, dx)
q30 = calculateOutgo(x30, u30, dx)
q40 = calculateOutgo(x40, u40, dx)
q50 = calculateOutgo(x50, u50, dx)
q60 = calculateOutgo(x60, u60, dx)
q70 = calculateOutgo(x70, u70, dx)
q80 = calculateOutgo(x80, u80, dx)
q90 = calculateOutgo(x90, u80, dx)

# Plot

plt.figure('Скорость и расход', figsize=(10, 7), dpi=200)

plt.title('Скорость потока воздуха\nв сечении затопленной струи')
plt.xlabel('Положение трубки Пито относительно центра струи [мм]')
plt.ylabel('Скоровть воздуха [м/с]')

plt.plot(x00 * 1000, u00, label='Q (00 мм) = {:.2f} [г/с]'.format(q00 * 1000))
plt.plot(x10 * 1000, u10, label='Q (10 мм) = {:.2f} [г/с]'.format(q10 * 1000))
plt.plot(x20 * 1000, u20, label='Q (20 мм) = {:.2f} [г/с]'.format(q20 * 1000))
plt.plot(x30 * 1000, u30, label='Q (30 мм) = {:.2f} [г/с]'.format(q30 * 1000))
plt.plot(x40 * 1000, u40, label='Q (40 мм) = {:.2f} [г/с]'.format(q40 * 1000))
plt.plot(x50 * 1000, u50, label='Q (50 мм) = {:.2f} [г/с]'.format(q50 * 1000))
plt.plot(x60 * 1000, u60, label='Q (60 мм) = {:.2f} [г/с]'.format(q60 * 1000))
plt.plot(x70 * 1000, u70, label='Q (70 мм) = {:.2f} [г/с]'.format(q70 * 1000))
plt.plot(x80 * 1000, u80, label='Q (80 мм) = {:.2f} [г/с]'.format(q80 * 1000))
plt.plot(x90 * 1000, u90, label='Q (90 мм) = {:.2f} [г/с]'.format(q90 * 1000))
plt.legend()

plt.grid(which='major')
plt.grid(which='minor', linestyle='--')
plt.minorticks_on()

plt.savefig(plotsPath + 'velocity-outgo.png')