import argparse
import sys
from bodystat import Bodystat
from inbody import Inbody
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
import numpy as np
import time

bs_path = "data/bodystat/ok/"
ib_path = "data/inbody/ok/"
bs_out_path = "data/bodystat/out/"
ib_out_path = "data/inbody/out/"

parser = argparse.ArgumentParser(description='')
parser.add_argument('-i', '--inbody', help='InBody only', required=False, action='store_true')
parser.add_argument('-b', '--bodystat', help='Bodystat only', required=False, action='store_true')
parser.add_argument('-c', '--compare', help='Compare Bodystat and Inbody device', required=False, action='store_true')
args = parser.parse_args()


def runBS():
    print(40 * '=' + 'BODYSTAT' + 40 * '=')
    bs = Bodystat(bs_path, bs_out_path)
    bs.exportData()
    bs.calculateNormalDist()
    bs.runStatisticTests()
    #bs.generateGraphs()
    #bs.getMeans()


def runIB():
    print(40 * '=' + 'INBODY' + 40 * '=')
    ib = Inbody(ib_path, ib_out_path)
    ib.exportData()
    ib.calculateNormalDist()
    ib.runStatisticTests()
    # ib.getMeans()


def pearsonCorr(k, l, bs_mes, ib_mes):
    print(l)
    r, p = pearsonr(bs_mes, ib_mes)
    print(f"Pearson correlation coeficient: {r.__round__(3)}\nP-value: {p.__round__(5)}")
    print('')

    plt.figure()
    plt.title(k + l)
    plt.boxplot([bs_mes, ib_mes])
    plt.xticks([1, 2], ["Bodystat QuadScan 4000", "InBody 770"])
    plt.savefig('data/both/boxplots/' + k + l)
    plt.close()
    plt.figure()

    bs_mes = np.array(bs_mes)
    ib_mes = np.array(ib_mes)
    plt.scatter(bs_mes, ib_mes)
    plt.xlabel('Bodystat QuadScan 4000')
    plt.ylabel('InBody 770')
    a = np.polyfit(bs_mes, ib_mes, 1)
    plt.plot(bs_mes, a[0]*bs_mes+a[1])
    plt.savefig('data/both/correlations/' + k + l)
    plt.close()


if args.compare:
    if args.bodystat and args.inbody:
        cpr_time = time.time()
        bs = Bodystat(bs_path, bs_out_path)
        ib = Inbody(ib_path, ib_out_path)
        mesList = bs.measurementList
        for j, k in enumerate(mesList):
            print(20 * '=' + k + 20 * '=')
            for l in ['_5khz', '_50khz', 'fat', 'water']:
                bs_mes = bs.getDataFromAll(data=l, mes=k)
                ib_mes = ib.getDataFromAll(data=l, mes=k)
                pearsonCorr(k, l, bs_mes, ib_mes)
        print(f"Compare execution time: {time.time() - cpr_time} ")
    else:
        print('You need to select both \'-i\' and \'-b\' flags for comparison!', file=sys.stderr)
        exit(1)
    exit(0)


if args.bodystat:
    bs_time = time.time()
    runBS()
    print(f"Bodystat execution time: {time.time() - bs_time} ")

if args.inbody:
    ib_time = time.time()
    runIB()
    print(f"InBody execution time: {time.time() - ib_time} ")



