main.py

from math import sqrt

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats


def clear_nans(no_nans, nans):
    """
    Removes elements from both lists at the positions where nans contains a NaN value
    :param no_nans: NaNs are ignored here. Gets elements removed on the basis of *nans* param
    :param nans: Contains NaNs whose position determine which elements are removed
    :return: *no_nans*, *nans*, cleared of elements at positions where *nans* contains a NaN value
    """
    assert (len(no_nans) == len(nans)), "Both lists must be the same length!"
    # WARNING: Do NOT attempt to change the order of filtering the arrays!
    no_nans = np.array([x for i, x in enumerate(no_nans) if not np.isnan(nans[i])])
    nans = np.array([x for i, x in enumerate(nans) if not np.isnan(nans[i])])
    return no_nans, nans


def plot_segment(slope, intercept, plot_range):
    """
    Returns x, y values for plotting a segment with given slope and y-intercept
    :param slope: Slope of segment
    :param intercept: y-intercept of segment
    :param plot_range: Range on which the segment should be plotted
    :return: Lists of *x* and *y* values
            which when plotted create a segment with given slope and y-intercept within the given range
    """
    x_v = np.array(plot_range)
    y_v = np.array([(x * slope + intercept) for x in x_v])
    return x_v, y_v


def fit_line(x, y):
    """
    Returns a slope and y-intercept of a fit line for values from x, y
    :param x: List of X values
    :param y: List of Y values
    :return:
    """
    x, y = clear_nans(x.values, y.values)
    slope, intercept, _, _, _ = stats.linregress(x, y)
    return slope, intercept


def plot_subplot(x, y_name, y, plot_range, subplot):
    slope, intercept = fit_line(x, y)
    x_line, y_line = plot_segment(slope, intercept, plot_range)
    subplot.plot(x_line, y_line)
    subplot.scatter(x, y, s=2)
    subplot.set_title(y_name)


def plot(name, filename, poundage, arrows):
    fig, axes = plt.subplots(2, 2, figsize=(8, 6))
    fig.suptitle(name, fontsize=15)

    for arrow in arrows:
        name, values, subplot_x, subplot_y = arrow
        plot_subplot(poundage, name, values, range(0, 200), axes[subplot_x, subplot_y])

    plt.savefig(filename, dpi=300)


def get_dispersion(x, y, slope, intercept):
    if np.isnan(y):
        return np.nan
    return -(slope * x - y + intercept) / sqrt(slope ** 2 + 1)


def compute_dispersion_list(poundage, distance):
    result = []

    slope, intercept = fit_line(poundage, distance)

    for el in zip(poundage, distance):
        poundage, distance = el
        dispersion = get_dispersion(poundage, distance, slope, intercept)
        result.append(dispersion)
    return result


if __name__ == "__main__":
    data = pd.read_csv('data.csv', delimiter=';')

    archer = data['Strelec']
    bowyer = data['Lukár']
    wood = data['Drevo']
    event = data['Akcia']

    poundage = data['Sila luku']
    blbs = data['BLBS']
    livery = data['Livery']
    fourth = data['A fourth pound']
    flight = data['Flight']

    plot("Datapoints", "datapoints.png", poundage,
         [("BLBS", blbs, 0, 0),
          ("Livery", livery, 0, 1),
          ("¼ pound", fourth, 1, 0),
          ("Flight", flight, 1, 1)])

    dispersion_blbs = compute_dispersion_list(poundage, blbs)
    dispersion_livery = compute_dispersion_list(poundage, livery)
    dispersion_fourth = compute_dispersion_list(poundage, fourth)
    dispersion_flight = compute_dispersion_list(poundage, flight)

    dispersion = pd.DataFrame(
        data={
            'Akcia': event,
            'Strelec': archer,
            'Lukár': bowyer,
            'Drevo': wood,
            'Sila luku': poundage,
            'BLBS': blbs,
            'BLBS disperzia': dispersion_blbs,
            'Livery': livery,
            'Livery disperzia': dispersion_livery,
            '¼ pound': fourth,
            '¼ pound disperzia': dispersion_fourth,
            'Flight': flight,
            'Flight disperzia': dispersion_flight
        },
        columns=['Akcia', 'Strelec', 'Lukár', 'Drevo', 'Sila luku',
                 'BLBS', 'BLBS disperzia',
                 'Livery', 'Livery disperzia',
                 '¼ pound', '¼ pound disperzia',
                 'Flight', 'Flight disperzia'])

    dispersion.to_csv('dispersion.csv', sep=';', encoding='utf-8')