import base64
from io import BytesIO
import matplotlib.pyplot as plt
import seaborn as sns
import re
from numpy import mean, std
import random
from sklearn.linear_model import LinearRegression
from sklearn.impute import SimpleImputer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.base import clone
import scipy.stats as stats
from .models import *
from .forms import *
import numpy as np
import pandas as pd
from scipy.stats import pearsonr

# function to get the Image and plot

def get_image() -> object:
    # create a byte buffer for the image to save
    buffer = BytesIO()
    # create a plot with the use of bytesio object
    plt.savefig(buffer, format='png')
    # set the cursor to the beginning of the screen
    buffer.seek(0)
    # retrieve the entire content of the file
    image_png = buffer.getvalue()
    # encoding and decoding
    graph = base64.b64encode(image_png)
    graph = graph.decode('utf-8')
    # free the memory of the buffer
    buffer.close()
    return graph

# enrollment capacity by school trend
def get_plot(chart_type, **kwargs):
    plt.switch_backend('AGG')
    # academic year
    x = kwargs.get('x')
    # y = total enrollment
    y = kwargs.get('y')
    # z = capacity of school
    z = kwargs.get('z')
    school_name = kwargs.get('name_of_school')

    if chart_type == 'barplot':
        fig, ax = plt.subplots(figsize=(10, 8))
        ax.bar(x, z, width=0.50, color='gray', label='capacity of school')
        ax.bar(x, y, width=0.25, color='g', label='total enrollment')
        ax.set_ylabel("Total Enrollment / Capacity")
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize(10)
            # specify integer or one of preset strings, e.g.
            # tick.label.set_fontsize('x-small')
            tick.label.set_rotation('60')
        for bar in ax.patches:
            # Using Matplotlib annotate function and
            # passing the coordinates where the annotation shall be done
            # x-coordinate: bar.get_x() + bar.get_width() / 2
            # y-coordinate: bar.get_height()
            # free space to be left to make graph pleasing: (0, 8)
            # ha and va stand for the horizontal and vertical alignment
            plt.annotate(format(bar.get_height(), '.2f'),
                         (bar.get_x() + bar.get_width() / 2,
                          bar.get_height()), ha='center', va='center',
                         size=8, xytext=(0, 3),
                         textcoords='offset points')
            plt.title(school_name)
        ax.set_ylim([0, max(z) + 300])
        ax.legend()

    else:
        title = school_name
        plt.figure(figsize=(10, 8))
        plt.title(title)
        plt.plot(x, y, 'b-', label='enrollments')
        plt.plot(x, z, 'g-', label='capacity')
        # plt.ylim(0, max(y) + 100)
        plt.ylabel("Total Enrollment / Capacity")
        plt.xlabel("Academic Year")
        plt.legend()
    plt.tight_layout()
    graph = get_image()
    return graph

# This graph plots trends of enrollment versus school capacity  for schools in a district for a select academic year
def get_pairs(**kwargs):
    plt.switch_backend('AGG')
    # total enrollment
    x = kwargs.get('x')
    # y = capacity of school
    y = kwargs.get('y')
    # selected year
    year = kwargs.get('academic_year')
    school = kwargs.get('name_of_school')
    district = kwargs.get('district_name')
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.bar(school, y, width=0.50, color='gray', label='capacity of school')
    ax.bar(school, x, width=0.25, color='g', label='total enrollment')
    ax.set_ylabel("Total Enrollment / Capacity")
    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        # specify integer or one of preset strings, e.g.
        # tick.label.set_fontsize('x-small')
        tick.label.set_rotation(45)
    for bar in ax.patches:
        # Using Matplotlib annotate function and
        # passing the coordinates where the annotation shall be done
        # x-coordinate: bar.get_x() + bar.get_width() / 2
        # y-coordinate: bar.get_height()
        # free space to be left to make graph pleasing: (0, 8)
        # ha and va stand for the horizontal and vertical alignment
        plt.annotate(format(bar.get_height(), '.2f'),
                     (bar.get_x() + bar.get_width() / 2,
                      bar.get_height()), ha='center', va='center',
                     size=8, xytext=(0, 3),
                     textcoords='offset points')
        plt.title("Enrollment/capacity for the academic year " + year + " of district  " + district)

    plt.legend()

    plt.tight_layout()
    graph = get_image()
    return graph

def get_grade_plot(**kwargs):
    plt.switch_backend('AGG')
    # the  selected academic year
    year = kwargs.get('academic_year')

    # the selected school
    school_name = kwargs.get('name_of_school')
    data_boys = kwargs.get('data_boys')
    data_girls = kwargs.get('data_girls')
    data_none = kwargs.get('data')

    # get subplots for boys, girls and single gender

    if len(data_none) > 0:
        fig, ax1 = plt.subplots(1, figsize=(10, 8))
        ax1.plot(data_none['grade'], data_none['enrollment'], 'ro')
        ax1.set_title('Enrollment')
        school_title = School.objects.filter(id=school_name).distinct().values_list('school_name', flat=True)
        school_name_title = school_title[0]
        fig.suptitle("Enrollment for " + school_name_title + " - Academic Year " + year)
        plt.xlabel('Grades')
        plt.tight_layout()
        ax1.grid()
        graph = get_image()
        return graph

    elif len(data_girls) > 0 and len(data_boys) > 0:
        fig, (ax1, ax2) = plt.subplots(2, figsize=(10, 8))
        ax1.plot(data_girls['grade'], data_girls['enrollment'], 'go')
        ax1.set_title('Girls Enrollment')
        ax2.plot(data_boys['grade'], data_boys['enrollment'], 'bo')
        ax2.set_title('Boys Enrollment')

        school_title = School.objects.filter(id=school_name).distinct().values_list('school_name', flat=True)
        school_name_title = school_title[0]
        fig.suptitle("Enrollment for " + school_name_title + " - Academic Year " + year)
        plt.xlabel('Grades')
        plt.tight_layout()
        ax1.grid()
        ax2.grid()
        graph = get_image()
        return graph
    else:
        graph = "No data recorded for the selected year"
        return graph

def get_district_grade_plot_girls(**kwargs):
    plt.switch_backend('AGG')
    # the  selected academic year
    year = kwargs.get('academic_year')
    grade = kwargs.get('grade')
    # the selected school
    district_name = kwargs.get('district')
    data_girls = kwargs.get('data_girls')
    school_girls = kwargs.get('school_girls')

    # get subplots for boys, girls and single gender
    fig, ax1 = plt.subplots(1, figsize=(10, 8))
    ax1.plot(data_girls['enrollment'], school_girls, 'go')
    ax1.set_title('Enrollment')
    plt.xticks(np.arange(0, max(data_girls['enrollment']) + 5, 5.0))
    for tick in ax1.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        # specify integer or one of preset strings, e.g.
        # tick.label.set_fontsize('x-small')
        tick.label.set_rotation(60)
    district_title = School.objects.filter(id=district_name).distinct().values_list('district_name', flat=True)
    district_name_title = district_title[0]
    fig.suptitle("Grade - " + str(grade) + " Girls Enrollment for Academic Year " + year + " for district" + str(
        district_name_title))
    plt.xlabel('Enrollment')
    plt.ylabel('School')
    plt.tight_layout()
    ax1.grid()
    graph = get_image()
    return graph

def get_district_grade_plot_boys(**kwargs):
    plt.switch_backend('AGG')
    # the  selected academic year
    year = kwargs.get('academic_year')
    grade = kwargs.get('grade')
    # the selected school
    district_name = kwargs.get('district')
    data_boys = kwargs.get('data_boys')
    school_boys = kwargs.get('school_boys')

    # get subplots for boys, girls and single gender
    fig, ax1 = plt.subplots(1, figsize=(10, 8))
    ax1.plot(data_boys['enrollment'], school_boys, 'bo')
    ax1.set_title('Enrollment')
    for tick in ax1.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        # specify integer or one of preset strings, e.g.
        # tick.label.set_fontsize('x-small')
        tick.label.set_rotation(60)
        plt.xticks(np.arange(0, max(data_boys['enrollment']) + 5, 5.0))
    district_title = School.objects.filter(id=district_name).distinct().values_list('district_name', flat=True)
    district_name_title = district_title[0]
    fig.suptitle("Grade - " + str(grade) + " Boys Enrollment for Academic Year " + year + " for district" + str(
        district_name_title))
    plt.xlabel('Enrollment')
    plt.ylabel('School')
    plt.tight_layout()
    ax1.grid()
    graph = get_image()
    return graph

def get_district_grade_plot_none(**kwargs):
    plt.switch_backend('AGG')
    # the  selected academic year
    year = kwargs.get('academic_year')
    grade = kwargs.get('grade')
    # the selected school
    district_name = kwargs.get('district')
    data_none = kwargs.get('data_none')
    school_none = kwargs.get('school_none')

    # get subplots for boys, girls and single gender
    fig, ax1 = plt.subplots(1, figsize=(10, 8))
    ax1.plot(data_none['enrollment'], school_none, 'ro')
    ax1.set_title('Enrollment')
    plt.xticks(np.arange(0, max(data_none['enrollment']) + 5, 5.0))
    for tick in ax1.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        # specify integer or one of preset strings, e.g.
        # tick.label.set_fontsize('x-small')
        tick.label.set_rotation(60)
    district_title = School.objects.filter(id=district_name).distinct().values_list('district_name', flat=True)
    district_name_title = district_title[0]
    fig.suptitle("Grade - " + str(grade) + " Total Enrollment for Academic Year " + year + " for district" + str(
        district_name_title))
    plt.xlabel('Enrollment')
    plt.ylabel('School')
    plt.tight_layout()
    ax1.grid()
    graph = get_image()
    return graph

def get_plot_boys_primary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    for _ in data.shape:
        ger = (data['enrollment'] / data['age_5_to_11_years']) * 100
        academic_year = data.academic_year
    sns.set(font_scale=1)
    sns.set_style("white")
    ax = ger.plot.bar(figsize=(15, 6))
    sns.despine(left=True, bottom=True)
    # label and title
    ax.set_xticklabels(np.arange(len(academic_year)))
    ax.set_title('Gross Enrollment Ratio (%) for Boys in Primary School In St. Lucia', size=18)
    ax.set_xticklabels(academic_year)
    for tick in ax.get_xticklabels():
        tick.set_rotation(-30)
    ax.set(xlabel='Academic Year', ylabel='Gross enrollment rate (%)')

    # annotations
    for p in ax.patches:
        ax.annotate(format(p.get_height(), '.2f'),
                    (p.get_x() + p.get_width() / 2., p.get_height()),
                    ha='center', va='center',
                    xytext=(0, 9),
                    textcoords='offset points')
    # adjust legend

    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_girls_primary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    for _ in data.shape:
        ger = (data['enrollment'] / data['age_5_to_11_years']) * 100
        academic_year = data.academic_year
    sns.set(font_scale=1)
    sns.set_style("white")
    ax = ger.plot.bar(figsize=(15, 6), color='green')
    sns.despine(left=True, bottom=True)
    # label and title
    ax.set_xticklabels(np.arange(len(academic_year)))
    ax.set_title('Gross Enrollment Ratio (%) for Girls in Primary School In St. Lucia', size=18)
    ax.set_xticklabels(academic_year)
    for tick in ax.get_xticklabels():
        tick.set_rotation(-30)
    ax.set(xlabel='Academic Year', ylabel='Gross enrollment rate (%)')

    # annotations
    for p in ax.patches:
        ax.annotate(format(p.get_height(), '.2f'),
                    (p.get_x() + p.get_width() / 2., p.get_height()),
                    ha='center', va='center',
                    xytext=(0, 9),
                    textcoords='offset points')
    # adjust legend

    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_primary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    data_boys = kwargs.get('data_boys')
    data_girls = kwargs.get('data_girls')

    for _ in data.shape:
        ger_boys = (data_boys['enrollment'] / data_boys['age_5_to_11_years']) * 100
        ger_girls = (data_girls['enrollment'] / data_girls['age_5_to_11_years']) * 100
        academic_year = data_girls.academic_year
        title = 'Trend of GER for Primary Schools in St. Lucia'
        plt.figure(figsize=(10, 8))
        plt.title(title)
        plt.plot(academic_year, ger_boys, 'b-', label='boys')
        plt.plot(academic_year, ger_girls, 'g-', label='girls')
        plt.xticks(rotation=60)
        # plt.ylim(0, max(y) + 100)
        plt.ylabel("Gross Enrollment Ratio for boys and girls in St. Lucia")
        plt.xlabel("Academic Year")
        plt.legend()
        plt.grid()
    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_boys_secondary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    for _ in data.shape:
        ger = (data['enrollment'] / data['age_12_to_16_years']) * 100
        academic_year = data.academic_year
    sns.set(font_scale=1)
    sns.set_style("white")
    ax = ger.plot.bar(figsize=(15, 6))
    sns.despine(left=True, bottom=True)
    # label and title
    ax.set_xticklabels(np.arange(len(academic_year)))
    ax.set_title('Gross Enrollment Ratio (%) for Boys in Secondary School In St. Lucia', size=18)
    ax.set_xticklabels(academic_year)
    for tick in ax.get_xticklabels():
        tick.set_rotation(-30)
    ax.set(xlabel='Academic Year', ylabel='Gross enrollment rate (%)')

    # annotations
    for p in ax.patches:
        ax.annotate(format(p.get_height(), '.2f'),
                    (p.get_x() + p.get_width() / 2., p.get_height()),
                    ha='center', va='center',
                    xytext=(0, 9),
                    textcoords='offset points')
    # adjust legend

    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_girls_secondary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    for _ in data.shape:
        ger = (data['enrollment'] / data['age_12_to_16_years']) * 100
        academic_year = data.academic_year
    sns.set(font_scale=1)
    sns.set_style("white")
    ax = ger.plot.bar(figsize=(15, 6), color='green')
    sns.despine(left=True, bottom=True)
    # label and title
    ax.set_xticklabels(np.arange(len(academic_year)))
    ax.set_title('Gross Enrollment Ratio (%) for Girls in secondary School In St. Lucia', size=18)
    ax.set_xticklabels(academic_year)
    for tick in ax.get_xticklabels():
        tick.set_rotation(-30)
    ax.set(xlabel='Academic Year', ylabel='Gross enrollment rate (%)')

    # annotations
    for p in ax.patches:
        ax.annotate(format(p.get_height(), '.2f'),
                    (p.get_x() + p.get_width() / 2., p.get_height()),
                    ha='center', va='center',
                    xytext=(0, 9),
                    textcoords='offset points')
    # adjust legend

    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_secondary(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    data_boys = kwargs.get('data_boys')
    data_girls = kwargs.get('data_girls')

    for _ in data.shape:
        ger_boys = (data_boys['enrollment'] / data_boys['age_12_to_16_years']) * 100
        ger_girls = (data_girls['enrollment'] / data_girls['age_12_to_16_years']) * 100
        academic_year = data_girls.academic_year
        title = 'Trend of GER for Primary Schools in St. Lucia'
        plt.figure(figsize=(10, 8))
        plt.title(title)
        plt.plot(academic_year, ger_boys, 'b-', label='boys')
        plt.plot(academic_year, ger_girls, 'g-', label='girls')
        plt.xticks(rotation=60)
        # plt.ylim(0, max(y) + 100)
        plt.ylabel("Gross Enrollment Ratio for boys and girls in Secondary School")
        plt.xlabel("Academic Year")
        plt.legend()
        plt.grid()

    plt.tight_layout()
    graph = get_image()
    return graph

def clean_secondary_name(name):
    name = re.sub("[^a-zA-Z]+", "", name)
    name = name.lower().replace('secondary', "")
    name = name.replace('school', "")
    return ' '.join(name.split())

def match_name(name, schools, district_dict):
    for school in schools:
        if clean_secondary_name(name) == clean_secondary_name(getattr(school, 'school_name')):
            district_code = getattr(school, 'district_name_id')
            district_dict[name] = district_code
            return district_code
    return None

def get_district(school_code, schools, district_dict):
    for school in schools:
        if int(getattr(school, 'school_code')) == school_code:
            district = getattr(school, 'district_name_id')
            district_dict[school_code] = district
            return district
    return None

def csec_performance_plot(data, district_1, district_2):
    left_out = set()

    df = pd.DataFrame(data.values())
    plt.switch_backend('AGG')
    years = [int(y) for y in df['year'].drop_duplicates()]
    years.sort()
    min_year = min(years)

    schools = School.objects.all()
    # schools = School.objects.filter(category_of_school='public secondary')

    N_DISTRICTS = District.objects.count()
    scores = np.zeros((len(years), N_DISTRICTS))
    n_tests = np.zeros((len(years), N_DISTRICTS))
    passing_scores = np.zeros((len(years), N_DISTRICTS))

    # cache school to district matches
    district_dict = {}

    for index, row in df.iterrows():
        school_code = int(row['school_id'])
        if school_code in district_dict:
            district = district_dict[school_code]
        else:
            district = get_district(school_code, schools, district_dict)
        if not district:
            left_out.add(row['school_id'])
            continue
        year = int(row['year']) - min_year
        n_tests[year][district - 1] += 1
        score = row['overall_grade']
        if score == 'I' or score == 'II' or score == 'III':
            scores[year][district - 1] += 1

    passing_scores = 100 * scores / n_tests
    passing_scores = pd.DataFrame(passing_scores)

    labels = ['District ' + str(d + 1) for d in range(N_DISTRICTS)]
    if not (district_1 and district_2):
        for d in range(N_DISTRICTS):
            plt.plot(years, passing_scores[d])
    else:
        plt.plot(years, passing_scores[district_1 - 1])
        plt.plot(years, passing_scores[district_2 - 1])
        labels = ['District ' + str(district_1), 'District ' + str(district_2)]
    plt.xticks([min(years), max(years)])
    plt.legend(labels, loc='upper left', bbox_to_anchor=(1, 1.05))
    plt.title("Percentage of Passing Scores (CSEC)")
    plt.tight_layout()
    graph = get_image()

    plt.clf()
    passing_scores = passing_scores.T
    passing_scores.columns = years
    passing_scores.index = ['District ' + str(d + 1) for d in range(N_DISTRICTS)]
    ax = sns.heatmap(passing_scores, annot=True)
    plt.tight_layout()
    heatmap = get_image()
    return [graph, heatmap, left_out]

# ===================================================================
# Outlier detection at district level
# ===================================================================
def get_outlier_district_plot(**kwargs):
    plt.switch_backend('AGG')

    school_enrollment = kwargs.get('x')
    school_name = kwargs.get('y')

    datamean = kwargs.get('data_mean')
    input_school_type = kwargs.get('input_school_type')
    academic_year = kwargs.get('academic_year')
    district_input = kwargs.get('input_district')

    fig, ax1 = plt.subplots(figsize=(10, 8))

    ax1.set_title('Enrollment for District')
    ax1.set_xlabel('School_Name')
    ax1.set_ylabel('School_Scores')

    ax1.bar(school_name, school_enrollment, color='b')

    for tick in ax1.xaxis.get_major_ticks():
        tick.label.set_fontsize(8)
        tick.label.set_rotation('15')
    plt.plot(school_name, datamean, linewidth=5, ls='solid', color='r')

    plt.xlabel("School Name")
    plt.ylabel("Enrollment")

    plt.title(
        "Enrollment for " + input_school_type + " schools for district " + district_input + " and " + academic_year + " academic year ")

    plt.tight_layout()
    graph = get_image()
    return graph

    plt.tight_layout()
    graph = get_image()
    return graph

# ==========================================================================
# Outlier detection at national level
# ==========================================================================

def get_outlier_national_plot(**kwargs):
    plt.switch_backend('AGG')

    school_enrollment = kwargs.get('x')
    school_name = kwargs.get('y')

    datamean = kwargs.get('data_mean')
    input_school_type = kwargs.get('input_school_type')
    academic_year = kwargs.get('academic_year')

    fig, ax1 = plt.subplots(figsize=(12, 10))

    ax1.set_title('Enrollment for Selected Year')
    ax1.set_xlabel('School_Name')

    ax1.bar(school_name, school_enrollment, width=0.1, color='b')
    for tick in ax1.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        tick.label.set_rotation('vertical')
    plt.plot(school_name, datamean, linewidth=3, ls='solid', color='r')

    plt.xlabel("School Name")
    plt.ylabel("Enrollment")
    plt.title("Enrollment for " + input_school_type + " schools for year " + academic_year)

    plt.tight_layout()
    graph = get_image()
    return graph

    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_regression(**kwargs):
    plt.switch_backend('AGG')
    plt.figure(figsize=(10, 8))

    data = kwargs.get('data')
    sns.set_theme(color_codes=True)
    sns.regplot(x=data.enrollment, y=data.gdp_millions, data=data, x_estimator=np.mean, label='GDP');
    sns.regplot(x=data.enrollment, y=data.educational_expenditure, data=data, x_estimator=np.mean,
                label='Educational Expenditure');
    sns.regplot(x=data.enrollment, y=data.government_expenditure, data=data, x_estimator=np.mean,
                label='Government Expenditure');
    plt.xlabel("Enrollment")
    plt.ylabel("Expenditure")
    plt.title("Linear Regression - Enrollment / GDP / Education / government expenditure")
    plt.legend()
    plt.tight_layout()
    graph = get_image()
    return graph

def get_plot_gdp_regress(**kwargs):
    plt.switch_backend('AGG')
    plt.figure(figsize=(10, 8))
    data = kwargs.get('data')
    sns.set_theme(color_codes=True)
    sns.jointplot(x=data.enrollment, y=data.gdp_millions, data=data, x_estimator=np.mean,
                  label='GDP', kind="reg");
    plt.xlabel("Enrollment")
    plt.ylabel("Expenditure")
    plt.tight_layout()
    graph = get_image()
    return graph

def get_enrollment_joint_pearsons(**kwargs):
    plt.switch_backend('AGG')
    plt.figure(figsize=(10, 8))
    data = kwargs.get('data')
    sns.set_theme(color_codes=True)
    import scipy.stats as stats
    graph = sns.jointplot(data=data, x=data.enrollment, y=data.gdp_millions)
    r, p = stats.pearsonr(x=data.enrollment, y=data.gdp_millions)
    phantom, = graph.ax_joint.plot([], [], linestyle="", alpha=0)
    # here graph is not a ax but a joint grid, so we access the axis through ax_joint method
    graph.ax_joint.legend([phantom], ['r={:f}, p={:f}'.format(r, p)])
    plt.tight_layout()
    graph = get_image()
    return graph

def get_enrollment_joint_spearman(**kwargs):
    plt.switch_backend('AGG')
    plt.figure(figsize=(10, 8))
    data = kwargs.get('data')
    sns.set_theme(color_codes=True)
    graph = sns.jointplot(data=data, x=data.enrollment, y=data.gdp_millions)
    r, p = stats.spearmanr(data.enrollment, data.gdp_millions)
    phantom, = graph.ax_joint.plot([], [], linestyle="", alpha=0)
    # here graph is not a ax but a joint grid, so we access the axis through ax_joint method

    graph.ax_joint.legend([phantom], ['r={:f}, p={:f}'.format(r, p)])
    plt.tight_layout()
    graph = get_image()
    return graph

def get_enrollment_multicollinearity(**kwargs):
    plt.figure(figsize=(10, 8))
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    data = data[["educational_expenditure", "gdp_millions", "government_expenditure", "primary_school_expenditure",
                 "secondary_school_expenditure", "enrollment", "age_5_to_11_years", "age_12_to_16_years"]]
    sns.set(style='white')
    corr = data.corr()
    mask = np.zeros_like(corr, dtype=np.bool)
    mask[np.triu_indices_from(mask)] = True
    f, ax = plt.subplots(figsize=(12, 10))
    cmap = sns.diverging_palette(220, 10, as_cmap=True)
    sns.heatmap(corr, mask=mask, cmap=cmap, vmax=.9, center=0, square=True, linewidths=.5, annot=True,
                cbar_kws={'shrink': .5});
    plt.tight_layout()
    graph = get_image()
    return graph

def get_kernel_density(**kwargs):
    plt.figure(figsize=(10, 8))
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    sns.kdeplot(data=data.enrollment)
    sns.despine()
    plt.tight_layout()
    graph = get_image()
    return graph

def plot_national_gender_enrollment(**kwargs):
    plt.switch_backend('AGG')
    data_boys_primary = kwargs.get('data_boys_primary')
    data_boys_secondary = kwargs.get('data_boys_secondary')
    data_girls_primary = kwargs.get('data_girls_primary')
    data_girls_secondary = kwargs.get('data_girls_secondary')
    title = 'Trend in enrollments over time'
    plt.figure(figsize=(10, 8))
    plt.title(title)
    plt.plot(data_boys_primary['academic_year'], data_boys_primary['enrollment'], 'b-',
             label='Boys enrolled in Primary School')
    plt.plot(data_boys_secondary['academic_year'], data_boys_secondary['enrollment'], 'bo',
             label='Boys enrolled in Secondary School')
    plt.plot(data_girls_primary['academic_year'], data_girls_primary['enrollment'], 'r-',
             label='Girls enrolled in Primary School')
    plt.plot(data_girls_secondary['academic_year'], data_girls_secondary['enrollment'], 'ro',
             label='Girls enrolled in secondary School')

    plt.xticks(rotation=60)
    # plt.ylim(0, max(y) + 100)
    plt.ylabel("National Enrollment Trends")
    plt.xlabel("Academic Year")
    plt.legend()
    plt.grid()

    plt.tight_layout()
    graph = get_image()
    return graph

def national_gender_enrollment_hist(**kwargs):
    data_boys_primary = kwargs.get('data_boys_primary')
    data_boys_secondary = kwargs.get('data_boys_secondary')
    data_girls_primary = kwargs.get('data_girls_primary')
    data_girls_secondary = kwargs.get('data_girls_secondary')

    # boys primary mean of distribution
    mu_boys_primary = mean(data_boys_primary.enrollment)
    mu_girls_primary = mean(data_girls_primary.enrollment)
    mu_boys_secondary = mean(data_boys_secondary.enrollment)
    mu_girls_secondary = mean(data_girls_secondary.enrollment)

    sigma_boys_primary = std(data_boys_primary.enrollment)
    sigma_girls_primary = std(data_girls_primary.enrollment)
    sigma_boys_secondary = std(data_boys_secondary.enrollment)
    sigma_girls_secondary = std(data_girls_secondary.enrollment)

    x_mu_boys_primary = mu_boys_primary + sigma_boys_primary * np.random.randn(437)
    x_mu_girls_primary = mu_girls_primary + sigma_girls_primary * np.random.randn(437)
    x_mu_boys_secondary = mu_boys_secondary + sigma_boys_secondary * np.random.randn(437)
    x_mu_girls_secondary = mu_girls_secondary + sigma_girls_secondary * np.random.randn(437)
    num_bins = 50

    # fig, ax = plt.subplots()

    fig, axs = plt.subplots(2, 2, figsize=(15, 15))

    # the histogram of the data
    n_boys_primary, bins_boys_primary, patches_boys_primary = axs[0, 0].hist(x_mu_boys_primary, num_bins, density=True)
    n_boys_secondary, bins_boys_secondary, patches_boys_secondary = axs[0, 1].hist(x_mu_boys_secondary, num_bins,
                                                                                density=True)
    n_girls_primary, bins_girls_primary, patches_girls_primary = axs[1, 0].hist(x_mu_girls_primary, num_bins,
                                                                                density=True)
    n_girls_secondary, bins_girls_secondary, patches_girls_secondary = axs[1, 1].hist(x_mu_girls_secondary, num_bins,
                                                                                density=True)
    # add a 'best fit' line
    y_boys_primary = ((1 / (np.sqrt(2 * np.pi) * sigma_boys_primary)) *
                      np.exp(-0.5 * (1 / sigma_boys_primary * (bins_boys_primary - mu_boys_primary)) ** 2))
    y_boys_secondary = ((1 / (np.sqrt(2 * np.pi) * sigma_boys_secondary)) *
                        np.exp(-0.5 * (1 / sigma_boys_secondary * (bins_boys_secondary - mu_boys_secondary)) ** 2))
    y_girls_primary = ((1 / (np.sqrt(2 * np.pi) * sigma_girls_primary)) *
                       np.exp(-0.5 * (1 / sigma_girls_primary * (bins_girls_primary - mu_girls_primary)) ** 2))

    y_girls_secondary = ((1 / (np.sqrt(2 * np.pi) * sigma_girls_secondary)) *
                         np.exp(-0.5 * (1 / sigma_girls_secondary * (bins_girls_secondary - mu_girls_secondary)) ** 2))

    for ax in axs.flat:
        ax.set(xlabel='Enrollment', ylabel='Probability Density')

    # Hide x labels and tick labels for top plots and y ticks for right plots.
    for ax in axs.flat:
        ax.label_outer()

    axs[0, 0].plot(bins_boys_primary, y_boys_primary, '--')
    axs[0, 0].set_title('Primary-Boys')

    axs[0, 1].plot(bins_boys_secondary, y_boys_secondary, '--')
    axs[0, 1].set_title('Secondary-Boys')

    axs[1, 0].plot(bins_girls_primary, y_girls_primary, '--')
    axs[1, 0].set_title('Primary-Girls')

    axs[1, 1].plot(bins_girls_secondary, y_girls_secondary, '--')
    axs[1, 1].set_title('Secondary-Girls')

    plt.tight_layout()
    graph = get_image()
    return graph

def plot_national_education_census(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')

    title = 'Education Census Over time'
    plt.figure(figsize=(10, 8))
    plt.title(title)
    plt.plot(data['academic_year'], data['age_3_to_4_years'], 'b-',
             label='Population of Age Group 3-4')
    plt.plot(data['academic_year'], data['age_5_to_11_years'], 'y-',
             label='Population of Age Group > 5 and less than 12')
    plt.plot(data['academic_year'], data['age_12_to_16_years'], 'r-',
             label='population of children above 12 years old')

    plt.xticks(rotation=60)
    # plt.ylim(0, max(y) + 100)
    plt.ylabel("Education Census Trends")
    plt.xlabel("Academic Year")
    plt.legend()
    plt.grid()

    plt.tight_layout()
    graph = get_image()
    return graph

def national_education_census_hist(**kwargs):
    data = kwargs.get('data')
    # boys primary mean of distribution
    mu_data_3_4 = mean(data.age_3_to_4_years)
    mu_data_5_11 = mean(data.age_5_to_11_years)
    mu_data_12_16 = mean(data.age_12_to_16_years)

    sigma_data_3_4 = std(data.age_3_to_4_years)
    sigma_data_5_11 = std(data.age_5_to_11_years)
    sigma_data_12_16 = std(data.age_12_to_16_years)

    x_mu_data_3_4 = mu_data_3_4 + sigma_data_3_4 * np.random.randn(437)
    x_mu_data_5_11 = mu_data_5_11 + sigma_data_5_11 * np.random.randn(437)
    x_mu_data_12_16 = mu_data_12_16 + sigma_data_12_16 * np.random.randn(437)
    num_bins = 50

    # fig, ax = plt.subplots()

    fig, (ax1, ax2, ax3) = plt.subplots(3, figsize=(15, 15))

    # the histogram of the data
    n_3_4, bins_3_4, patches_3_4 = ax1.hist(x_mu_data_3_4, num_bins, density=True)
    n_5_11, bins_5_11, patches_5_11 = ax2.hist(x_mu_data_5_11, num_bins, density=True)
    n_12_16, bins_12_16, patches_12_16 = ax3.hist(x_mu_data_12_16, num_bins, density=True)

    # add a 'best fit' line
    y_3_4 = ((1 / (np.sqrt(2 * np.pi) * sigma_data_3_4)) *
             np.exp(-0.5 * (1 / sigma_data_3_4 * (bins_3_4 - mu_data_3_4)) ** 2))
    y_5_11 = ((1 / (np.sqrt(2 * np.pi) * sigma_data_5_11)) *
              np.exp(-0.5 * (1 / sigma_data_5_11 * (bins_5_11 - mu_data_5_11)) ** 2))
    y_12_16 = ((1 / (np.sqrt(2 * np.pi) * sigma_data_12_16)) *
               np.exp(-0.5 * (1 / sigma_data_12_16 * (bins_12_16 - mu_data_12_16)) ** 2))

    ax1.plot(bins_3_4, y_3_4, '--')
    ax1.set_title('3 - 4 years')

    ax2.plot(bins_5_11, y_5_11, '--')
    ax2.set_title('Greater / equal 5, Less than 12')

    ax3.plot(bins_12_16, y_12_16, '--')
    ax3.set_title('>=12 years')

    plt.tight_layout()
    graph = get_image()
    return graph

def plot_national_expenditure(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    title = 'Education Expenditure'
    plt.figure(figsize=(10, 8))
    plt.title(title)
    plt.plot(data['academic_year'], data.educational_expenditure, 'b-',
             label='Educational Expenditure')
    plt.plot(data['academic_year'], data['gdp_millions'], 'y-',
             label='GDP (Million XCD)')
    plt.plot(data['academic_year'], data['government_expenditure'], 'r-',
             label='Government Expenditure')
    plt.plot(data['academic_year'], data['primary_school_expenditure'], 'g-',
             label='Primary School Expenditure')
    plt.plot(data['academic_year'], data['secondary_school_expenditure'], 'k-',
             label='Secondary School Expenditure')

    plt.xticks(rotation=60)
    # plt.ylim(0, max(y) + 100)
    plt.ylabel("Expenditure Trends")
    plt.xlabel("Academic Year")
    plt.legend()
    plt.grid()

    plt.tight_layout()
    graph = get_image()
    return graph

def national_expenditure_hist(**kwargs):
    data = kwargs.get('data')
    mu_educational_expenditure = mean(data['educational_expenditure'])
    mu_gdp_millions = mean(data['gdp_millions'])
    mu_government_expenditure = mean(data['government_expenditure'])
    mu_primary_school_expenditure = mean(data['primary_school_expenditure'])
    mu_secondary_school_expenditure = mean(data['secondary_school_expenditure'])

    sigma_educational_expenditure = std(data.educational_expenditure)
    sigma_gdp_millions = std(data.gdp_millions)
    sigma_government_expenditure = std(data.government_expenditure)
    sigma_primary_school_expenditure = std(data.primary_school_expenditure)
    sigma_secondary_school_expenditure = std(data.secondary_school_expenditure)

    x_mu_educational_expenditure = mu_educational_expenditure + sigma_educational_expenditure * np.random.randn(437)
    x_mu_gdp_millions = mu_gdp_millions + sigma_gdp_millions * np.random.randn(437)
    x_mu_government_expenditure = mu_government_expenditure + sigma_government_expenditure * np.random.randn(437)
    x_mu_primary_school_expenditure = mu_primary_school_expenditure + sigma_primary_school_expenditure * np.random.randn(
        437)
    x_mu_secondary_school_expenditure = mu_secondary_school_expenditure + sigma_secondary_school_expenditure * np.random.randn(
        437)
    num_bins = 50

    # fig, ax = plt.subplots()

    fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(5, figsize=(15, 15))

    # the histogram of the data
    n_educational_expenditure, bins_educational_expenditure, patches_educational_expenditure = \
        ax1.hist(x_mu_educational_expenditure, num_bins, density=True)
    n_gdp_millions, bins_gdp_millions, patches_gdp_millions = \
        ax2.hist(x_mu_gdp_millions, num_bins, density=True)
    n_government_expenditure, bins_government_expenditure, patches_government_expenditure = \
        ax3.hist(x_mu_government_expenditure, num_bins, density=True)
    n_primary_school_expenditure, bins_primary_school_expenditure, patches_primary_school_expenditure = \
        ax4.hist(x_mu_primary_school_expenditure, num_bins, density=True)
    n_secondary_school_expenditure, bins_secondary_school_expenditure, patches_secondary_school_expenditure = \
        ax5.hist(x_mu_secondary_school_expenditure, num_bins, density=True)
    # add a 'best fit' line
    y_educational_expenditure = ((1 / (np.sqrt(2 * np.pi) * sigma_educational_expenditure)) *
                                 np.exp(-0.5 * (1 / sigma_educational_expenditure * (
                                         bins_educational_expenditure - mu_educational_expenditure)) ** 2))
    y_gdp_millions = ((1 / (np.sqrt(2 * np.pi) * sigma_gdp_millions)) *
                      np.exp(-0.5 * (1 / sigma_gdp_millions * (bins_gdp_millions - mu_gdp_millions)) ** 2))
    y_government_expenditure = ((1 / (np.sqrt(2 * np.pi) * sigma_government_expenditure)) *
                                np.exp(-0.5 * (1 / sigma_government_expenditure * (
                                        bins_government_expenditure - mu_government_expenditure)) ** 2))
    y_primary_school_expenditure = ((1 / (np.sqrt(2 * np.pi) * sigma_primary_school_expenditure)) *
                                    np.exp(-0.5 * (1 / sigma_primary_school_expenditure * (
                                            bins_primary_school_expenditure - mu_primary_school_expenditure)) ** 2))
    y_secondary_school_expenditure = ((1 / (np.sqrt(2 * np.pi) * sigma_secondary_school_expenditure)) *
                                      np.exp(-0.5 * (1 / sigma_secondary_school_expenditure * (
                                              bins_secondary_school_expenditure - mu_secondary_school_expenditure)) ** 2))

    ax1.plot(bins_educational_expenditure, y_educational_expenditure, '--')
    ax1.set_title('Educational Expenditure')

    ax2.plot(bins_gdp_millions, y_gdp_millions, '--')
    ax2.set_title('gdp(Million XCD)')

    ax3.plot(bins_government_expenditure, y_government_expenditure, '--')
    ax3.set_title('Government Expenditure')

    ax4.plot(bins_primary_school_expenditure, y_primary_school_expenditure, '--')
    ax4.set_title('Primary School Expenditure')

    ax5.plot(bins_secondary_school_expenditure, y_secondary_school_expenditure, '--')
    ax5.set_title('Secondary School Expenditure')

    plt.tight_layout()
    graph = get_image()
    return graph

def primary_performance_plot(data, district_1, district_2):
    df = pd.DataFrame(data.values())
    plt.switch_backend('AGG')
    years = df['academic_year'].drop_duplicates().str.split("/")
    years = [int(y[1]) for y in years]
    min_year = min(years)

    N_DISTRICTS = District.objects.count()

    tests = np.zeros((len(years), N_DISTRICTS), dtype=int)
    above_avg = np.zeros((len(years), N_DISTRICTS), dtype=int)
    performance = np.zeros((len(years), N_DISTRICTS), dtype=float)

    for index, row in df.iterrows():
        year = int(row['academic_year'].split('/')[1])
        school_code = row['school_id']
        school = School.objects.get(school_code=school_code)
        district = getattr(school, 'district_name_id')
        n_tests = int(row['tests_sat'])
        n_above_avg = int(row['above_average_scores'])
        if np.isnan(n_tests) or np.isnan(n_above_avg):
            continue
        tests[year - min_year][district - 1] += n_tests
        above_avg[year - min_year][district - 1] += n_above_avg

    for y in range(len(years)):
        performance[y] = 100 * above_avg[y] / tests[y]
    performance = pd.DataFrame(performance)
    labels = ['District ' + str(d + 1) for d in range(N_DISTRICTS)]
    if not (district_1 and district_2):
        for d in range(N_DISTRICTS):
            plt.plot(years, performance[d])
    else:
        plt.plot(years, performance[district_1 - 1])
        plt.plot(years, performance[district_2 - 1])
        labels = ['District ' + str(district_1), 'District' + str(district_2)]
    plt.xticks([min(years), max(years)])
    plt.legend(labels, loc='upper left', bbox_to_anchor=(1, 1.05))
    plt.title("Percentage of Students Scoring Above Mean (CEE)")
    plt.tight_layout()
    graph = get_image()

    plt.clf()
    performance = performance.T
    performance.columns = np.arange(1999, 2018, step=1)
    performance.index = ['District ' + str(d + 1) for d in range(N_DISTRICTS)]
    ax = sns.heatmap(performance, annot=True)
    ax.set_title("Percentage of Students Scoring above Mean (CEE)")
    plt.tight_layout()
    heatmap = get_image()
    return [graph, heatmap]

def get_sex(character):
    if character == "F":
        return "female"
    else:
        return "male"

def store_scores(data, required_fields, user_data, type):
    result = {}
    lines = data.replace("\r", "").split("\n")
    field_names = lines[0].split(",")
    if not set(required_fields).issubset(set(field_names)):
        diff = set(required_fields) - set(field_names)
        missing_fields = []
        for d in diff:
            missing_fields.append(d)
        result['missing_fields'] = missing_fields
        result['error_message'] = 'The following fields are missing:\n'
        result['n_scores'] = 0
    else:
        succeeded = 0
        failed = 0
        for line in lines[1:]:
            if line:
                fields = line.split(",")
                data = {}
                for required_field in required_fields:
                    if required_field == "school_id":
                        data["school"] = fields[field_names.index("school_id")]
                    elif required_field == "primsch_id":
                        data["primsch"] = fields[field_names.index("primsch_id")]
                    elif required_field == "secsch_id":
                        data["secsch"] = fields[field_names.index("secsch_id")]
                    elif required_field == "district_id":
                        data["district"] = fields[field_names.index("district_id")]
                    elif required_field == "sex":
                        data["sex"] = get_sex(fields[field_names.index(required_field)])
                    else:
                        data[required_field] = fields[field_names.index(required_field)]
                data = {**data, **user_data}
                if type == "CEE":
                    form = CEEForm(data)
                if type == "CSEC":
                    form = CSECForm(data)
                if form.is_valid():
                    form.save()
                    succeeded += 1
                else:
                    failed += 1
        result['n_scores'] = succeeded
        result['failed'] = failed
    return result

# =======================================================================================
# Box plots at district level
# =======================================================================================

def get_boxplot_district_plot(**kwargs):
    plt.switch_backend('AGG')

    school_enrollment = kwargs.get('x')
    school_name = kwargs.get('y')

    input_school_type = kwargs.get('input_school_type')
    academic_year = kwargs.get('academic_year')
    district_input = kwargs.get('input_district')

    fig, ax1 = plt.subplots(figsize=(11, 6))

    plt.boxplot(school_enrollment, patch_artist=True,
                boxprops=dict(facecolor='purple'),
                meanline=True, showmeans=True)

    plt.xticks([1], [input_school_type])
    plt.ylabel('Enrollment')
    plt.title(
        "Box Plot for Enrollment in " + input_school_type + " schools  in " + " District" + district_input + " and academic year " + academic_year)

    plt.tight_layout()
    graph = get_image()
    return graph

# =========================================================================================
# Box plots at national level
# =========================================================================================

def get_boxplot_national_plot(**kwargs):
    plt.switch_backend('AGG')

    school_enrollment = kwargs.get('x')
    school_name = kwargs.get('y')

    input_school_type = kwargs.get('input_school_type')
    academic_year = kwargs.get('academic_year')

    fig, ax1 = plt.subplots(figsize=(11, 6))

    plt.boxplot(school_enrollment, patch_artist=True,
                boxprops=dict(facecolor='purple'),
                meanline=True, showmeans=True)

    plt.xticks([1], [input_school_type])
    plt.ylabel('Enrollment')
    plt.title("Box Plot for Enrollment in " + input_school_type + " schools " + " for academic year " + academic_year)

    plt.tight_layout()
    graph = get_image()
    return graph

def plot_national_ratio_trend(**kwargs):
    plt.switch_backend('AGG')
    data_primary = kwargs.get('data_primary')
    data_secondary = kwargs.get('data_secondary')
    for _ in data_secondary.shape:
        school_enrollment = data_secondary.total_enrollment
        total_number_of_teachers_secondary = data_secondary.total_number_of_teachers
        academic_year = data_secondary.academic_year
        student_teacher_ratio_secondary = (
            (school_enrollment / total_number_of_teachers_secondary).replace(np.inf, 0)).astype(float)
        fig, ax = plt.subplots(figsize=(12, 10))
        ax.bar(academic_year, student_teacher_ratio_secondary, width=0.8, color='b',
               label='Student to teacher ratio in Secondary Schools'
               , alpha=0.5)
        ax.set_ylabel("Student - Teacher Ratio (7-11)")

        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize(10)
            tick.label.set_rotation(45)
        for bar in ax.patches:
            plt.annotate(format(bar.get_height(), '.2f'),
                         (bar.get_x() + bar.get_width() / 2,
                          bar.get_height()), ha='center', va='center',
                         size=7, xytext=(0, 5),
                         textcoords='offset points')
            plt.title("Student to Teacher Ratio Trends in Secondary Schools")
            plt.xlabel('Academic Year')
            ax.set_ylim([0, max(student_teacher_ratio_secondary) + 100])

            ax.legend()

    plt.tight_layout()
    # plt.grid()
    graph = get_image()
    return graph

def plot_national_ratio_trend_primary(**kwargs):
    plt.switch_backend('AGG')
    data_primary = kwargs.get('data_primary')
    for _ in data_primary.shape:
        school_enrollment = data_primary.total_enrollment
        total_number_of_teachers_primary = data_primary.total_number_of_teachers
        academic_year = data_primary.academic_year
        student_teacher_ratio_primary = (
            (school_enrollment / total_number_of_teachers_primary).replace(np.inf, 0)).astype(float)
        fig, ax = plt.subplots(figsize=(12, 10))
        ax.bar(academic_year, student_teacher_ratio_primary, width=0.8, color='g',
               label='Student to teacher ratio in Primary Schools'
               , alpha=0.5)
        ax.set_ylabel("Student - Teacher Ratio (k-6)")

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(10)
        tick.label.set_rotation(45)
    for bar in ax.patches:
        plt.annotate(format(bar.get_height(), '.2f'),
                     (bar.get_x() + bar.get_width() / 2,
                      bar.get_height()), ha='center', va='center',
                     size=7, xytext=(0, 5),
                     textcoords='offset points')
        plt.title("Student to Teacher Ratio Trends in Primary Schools")
        plt.xlabel('Academic Year')
        ax.set_ylim([0, max(student_teacher_ratio_primary) + 100])

        ax.legend()
    plt.tight_layout()
    # plt.grid()
    graph = get_image()
    return graph

def national_ratio_hist(**kwargs):
    plt.switch_backend('AGG')
    data = kwargs.get('data')
    mu_total_enrollment = mean(data['total_enrollment'])
    mu_number_of_trained_male_teachers = mean(data['number_of_trained_male_teachers'])
    mu_number_of_trained_female_teachers = mean(data['number_of_trained_female_teachers'])
    mu_number_of_untrained_male_teachers = mean(data['number_of_untrained_male_teachers'])
    mu_number_of_untrained_female_teachers = mean(data['number_of_untrained_female_teachers'])
    mu_total_no_of_teachers = mean(data['total_number_of_teachers'])

    sigma_total_enrollment = std(data['total_enrollment'])
    sigma_number_of_trained_male_teachers = std(data['number_of_trained_male_teachers'])
    sigma_number_of_trained_female_teachers = std(data['number_of_trained_female_teachers'])
    sigma_number_of_untrained_male_teachers = std(data['number_of_untrained_male_teachers'])
    sigma_number_of_untrained_female_teachers = std(data['number_of_untrained_female_teachers'])
    sigma_total_no_of_teachers = std(data['total_number_of_teachers'])

    x_mu_total_enrollment = mu_total_enrollment + sigma_total_enrollment * np.random.randn(437)
    x_mu_number_of_trained_male_teachers = mu_number_of_trained_male_teachers + sigma_number_of_trained_male_teachers * np.random.randn(
        437)
    x_mu_number_of_trained_female_teachers = mu_number_of_trained_female_teachers + sigma_number_of_trained_female_teachers * np.random.randn(
        437)
    x_mu_number_of_untrained_male_teachers = mu_number_of_untrained_male_teachers + sigma_number_of_untrained_male_teachers * np.random.randn(
        437)
    x_mu_number_of_untrained_female_teachers = mu_number_of_untrained_female_teachers + sigma_number_of_untrained_female_teachers * np.random.randn(
        437)
    x_mu_total_no_of_teachers = mu_total_no_of_teachers + sigma_total_no_of_teachers * np.random.randn(
        437)
    num_bins = 50

    # fig, ax = plt.subplots()

    fig, (ax1, ax2, ax3, ax4, ax5, ax6) = plt.subplots(6, figsize=(15, 15))

    # the histogram of the data
    n_total_enrollment, bins_total_enrollment, patches_total_enrollment = \
        ax1.hist(x_mu_total_enrollment, num_bins, density=True)
    n_number_of_trained_male_teachers, bins_number_of_trained_male_teachers, patches_number_of_trained_male_teachers = \
        ax2.hist(x_mu_number_of_trained_male_teachers, num_bins, density=True)
    n_number_of_trained_female_teachers, bins_number_of_trained_female_teachers, patches_number_of_trained_female_teachers = \
        ax3.hist(x_mu_number_of_trained_female_teachers, num_bins, density=True)
    n_number_of_untrained_male_teachers, bins_number_of_untrained_male_teachers, patches_number_of_untrained_male_teachers = \
        ax4.hist(x_mu_number_of_untrained_male_teachers, num_bins, density=True)
    n_number_of_untrained_female_teachers, bins_number_of_untrained_female_teachers, patches_number_of_untrained_female_teachers = \
        ax5.hist(x_mu_number_of_untrained_female_teachers, num_bins, density=True)
    n_total_number_of_teachers, bins_total_number_of_teachers, patches_total_number_of_teachers = \
        ax6.hist(x_mu_total_no_of_teachers, num_bins, density=True)
    # add a 'best fit' line
    y_total_enrollment = ((1 / (np.sqrt(2 * np.pi) * sigma_total_enrollment)) *
                          np.exp(-0.5 * (1 / sigma_total_enrollment * (
                                  bins_total_enrollment - mu_total_enrollment)) ** 2))
    y_number_of_trained_male_teachers = ((1 / (np.sqrt(2 * np.pi) * sigma_number_of_trained_male_teachers)) *
                                         np.exp(-0.5 * (1 / sigma_number_of_trained_male_teachers * (
                                                 bins_number_of_trained_male_teachers - mu_number_of_trained_male_teachers)) ** 2))
    y_number_of_trained_female_teachers = ((1 / (np.sqrt(2 * np.pi) * sigma_number_of_trained_female_teachers)) *
                                           np.exp(-0.5 * (1 / sigma_number_of_trained_female_teachers * (
                                                   bins_number_of_trained_female_teachers - mu_number_of_trained_female_teachers)) ** 2))
    y_number_of_untrained_male_teachers = ((1 / (np.sqrt(2 * np.pi) * sigma_number_of_untrained_male_teachers)) *
                                           np.exp(-0.5 * (1 / sigma_number_of_untrained_male_teachers * (
                                                   bins_number_of_untrained_male_teachers - mu_number_of_untrained_male_teachers)) ** 2))
    y_number_of_untrained_female_teachers = ((1 / (np.sqrt(2 * np.pi) * sigma_number_of_untrained_female_teachers)) *
                                             np.exp(-0.5 * (1 / sigma_number_of_untrained_female_teachers * (
                                                     bins_number_of_untrained_female_teachers - mu_number_of_untrained_female_teachers)) ** 2))
    y_total_no_of_teachers = ((1 / (np.sqrt(2 * np.pi) * sigma_total_no_of_teachers)) *
                              np.exp(-0.5 * (1 / sigma_total_no_of_teachers * (
                                      bins_total_number_of_teachers - mu_total_no_of_teachers)) ** 2))

    ax1.plot(bins_total_enrollment, y_total_enrollment, '--')
    ax1.set_title('Total Enrollment')

    ax2.plot(bins_number_of_trained_male_teachers, y_number_of_trained_male_teachers, '--')
    ax2.set_title('Number of Trained Male Teachers')

    ax3.plot(bins_number_of_trained_female_teachers, y_number_of_trained_female_teachers, '--')
    ax3.set_title('Number of Trained Female Teachers')

    ax4.plot(bins_number_of_untrained_male_teachers, y_number_of_untrained_male_teachers, '--')
    ax4.set_title('Number of Untrained Male Teachers')

    ax5.plot(bins_number_of_untrained_female_teachers, y_number_of_untrained_female_teachers, '--')
    ax5.set_title('Untrained Female Teachers')

    ax6.plot(bins_total_number_of_teachers, y_total_no_of_teachers, '--')
    ax6.set_title('Total Number of Teachers')

    plt.tight_layout()
    graph = get_image()
    return graph

def get_prev_yr(year_string):
    year = int(year_string.split("/")[0])
    prev_year = year - 1
    return str(prev_year) + "/" + str(year)

def get_previous_performance(data):
    previous = []
    for index, row in data.iterrows():
        previous_year = data.loc[(data['school_id'] == row['school_id']) &
                                 (data['academic_year'] == get_prev_yr(row['academic_year']))]
        if not previous_year.empty:
            previous.append(previous_year['performance'].values[0])
        else:
            previous.append(np.nan)

    previous = pd.Series(previous)
    data['previous'] = previous
    return data

def get_enrollment(data):
    enrollment = []
    capacity = []
    found = 0
    not_found = 0
    for index, row in data.iterrows():
        school = AggregateEnrollment.objects.filter(name_of_school_id=row['school_id'],
                                                    academic_year=row['academic_year'])
        if school:
            enr = getattr(school[0], 'total_enrollment')
            enrollment.append(enr)
            capacity.append(enr / getattr(school[0], 'capacity_of_school'))
        else:
            enrollment.append(np.nan)
            capacity.append(np.nan)

    enrollment = pd.Series(enrollment)
    data['enrollment'] = enrollment
    capacity = pd.Series(capacity)
    data['capacity'] = capacity
    return data

def divide_by_enrollment(data):
    data['teachers'] = data['enrollment'] / data['teachers']
    data['girls'] = data['girls'] / data['enrollment']
    data['total_bursaries'] = data['total_bursaries'] / data['enrollment']
    data['feeding_program'] = data['feeding_program'] / data['enrollment']
    data['repeaters'] = data['repeaters'] / data['enrollment']
    data['trained_teachers'] = data['trained_teachers'] / data['teachers']
    data['primary_enrollment'] = data['primary_enrollment'] / data['enrollment']
    return data

def correlations(data, excluded_fields):
    plt.switch_backend('AGG')
    data = pd.DataFrame(data.values())
    data['performance'] = data['above_average_scores'] / data['tests_sat']
    data = data.drop(columns=excluded_fields)

    data = get_previous_performance(data)
    data = get_enrollment(data)

    data = data.drop(columns=['academic_year', 'school_id', 'above_average_scores', 'tests_sat'])
    data = data.apply(pd.to_numeric)
    data = divide_by_enrollment(data)

    correlation = []
    spearman = []
    for f in data.columns:
        if f == 'performance':
            continue
        df = data[data[f].notna()]
        x = df[f]
        y = df['performance']
        correlation.append(stats.pearsonr(x, y))
        spearman.append(stats.spearmanr(x, y))
    correlation = pd.DataFrame(correlation)
    spearman = pd.DataFrame(spearman)
    correlation.columns = ['r (Pe)', 'p (Pe)']
    correlation.index = data.columns.drop('performance')
    spearman.index = data.columns.drop('performance')
    correlation['r (Sp)'] = spearman['correlation']
    correlation['p (Sp)'] = spearman['pvalue']

    mask = np.zeros((len(correlation), 4))
    mask[:, 3] = True
    mask[:, 1] = True
    ax = sns.heatmap(correlation, annot=True, mask=mask)
    ax.set_title("Correlations between School Factors and Exam Performance")

    for (j, i), label in np.ndenumerate(correlation.values):
        label = "{:.2e}".format(label)
        if i == 1 or i == 3:
            ax.text(i + 0.5, j + 0.5, label,
                    fontdict=dict(ha='center', va='center', color='black'))
    plt.tight_layout()
    graph = get_image()
    return graph

def dropcol_importances(rf, X_train, y_train):
    r = random.randint(1, 999)
    rf_ = clone(rf)
    rf_.random_state = r
    rf_.fit(X_train, y_train)
    baseline = rf_.oob_score_
    imp = []
    for col in X_train.columns:
        X = X_train.drop(col, axis=1)
        rf_ = clone(rf)
        rf_.random_state = r
        rf_.fit(X, y_train)
        o = rf_.oob_score_
        imp.append(baseline - o)
    imp = np.array(imp)
    I = pd.DataFrame(
        data={'Feature': X_train.columns,
              'Importance': imp})
    I = I.set_index('Feature')
    I = I.sort_values('Importance', ascending=True)
    return I

def rf_model(data, excluded_fields, random):
    plt.switch_backend('AGG')
    data = pd.DataFrame(data.values())
    data['performance'] = data['above_average_scores'] / data['tests_sat']
    data = data.drop(columns=excluded_fields)

    data = get_previous_performance(data)
    data = get_enrollment(data)
    data = data.drop(columns=['academic_year', 'school_id', 'above_average_scores', 'tests_sat'])

    data = data.apply(pd.to_numeric)
    data = divide_by_enrollment(data)

    imp_mean = SimpleImputer(missing_values=np.nan, strategy='median')
    imp_mean.fit(data)
    SimpleImputer()
    imputed_data = pd.DataFrame(imp_mean.transform(data))
    imputed_data.columns = data.columns
    imputed_data.index = data.index

    data = imputed_data

    performance = data['performance']
    data = data.drop(columns=['performance'])

    X_train, X_test, y_train, y_test = train_test_split(data, performance, train_size=0.8)

    clf = RandomForestRegressor(max_features=None, oob_score=True)
    clf.fit(X_train, y_train)

    accuracy = round(clf.score(X_test, y_test), 3)

    data = data.drop(columns=['previous'])
    if random:
        data['random'] = pd.Series([random.random() for x in range(len(data))])
    importance = dropcol_importances(clf, data, performance)

    # plot feature importance
    ax = pd.Series(importance['Importance']).plot(kind='barh')
    y_labels = importance.index
    ax.set_yticklabels(y_labels)
    plt.xlabel("Change in OOB Error")
    plt.title("Random Forest Regression Feature Importance")

    plt.tight_layout()
    graph = get_image()
    return [graph, accuracy]


<gh_stars>0
"""Models for estimating the aerodynamics of a 3D foil from its sections."""

from __future__ import annotations

import abc
from typing import TYPE_CHECKING, Protocol, runtime_checkable

import numpy as np
import scipy.optimize

from pfh.glidersim.util import cross3

if TYPE_CHECKING:
    from pfh.glidersim.foil import SimpleFoil

__all__ = [
    "FoilAerodynamics",
    "ConvergenceError",
    "Phillips",
]

def __dir__():
    return __all__

@runtime_checkable
class FoilAerodynamics(Protocol):
    """Interface for classes that implement a FoilAerodynamics model."""

    @abc.abstractmethod
    def __call__(self, ai, v_W2f, rho_air, **kwargs):
        """
        Estimate the forces and moments on a foil.

        Parameters
        ----------
        ai : array_like of float
            Airfoil indices. The shape must be able to broadcast to (K,), where
            `K` is the number of control points being used by the estimator.
        v_W2f : array_like of float [m/s]
            The velocity of the wind relative to the control points in foil frd
            coordinates. The shape must be able to broadcast to (K, 3), where
            `K` is the number of control points being used by the estimator.
        rho_air : array_like of float [kg/m^3]
            Air density
        """

    @abc.abstractmethod
    def r_CP2LE(self):
        """
        Compute the control points for the section aerodynamics.

        Returns
        -------
        ndarray of float, shape (K,3) [m]
            Control points relative to the central leading edge `LE`.
            Coordinates are in canopy frd, and `K` is the number of points
            being used by the estimation method.
        """

class ConvergenceError(RuntimeError):
    """The estimator failed to converge on a solution."""

class Phillips(FoilAerodynamics):
    """
    A non-linear numerical lifting-line method.

    Uses a set of spanwise bound vortices instead of a single, uniform lifting
    line. Unlike the Prandtl's classic lifting-line theory, this method allows
    for wing sweep and dihedral.

    Parameters
    ----------
    foil : SimpleFoil
        Defines the lifting-line and section coefficients.
    v_ref_mag : float [m/s]
        The reference solution airspeed
    alpha_ref : float [degrees]
        The reference solution angle of attack
    s_nodes : array_like of floats, shape (K+1,)
        Section indices of the `K + 1` section nodes (wing segment endpoints).
        The `K >= 1` aerodynamic control points are centered between the nodes.
        Two common point distributions are:

        * Linear: ``np.linspace(-1, 1, K + 1)``
        * Cosine: ``np.cos(np.linspace(np.pi, 0, K + 1))``

    s_clamp : float, optional
        Section index to enable clamped output of the aerodynamic coefficients
        for section indices `abs(s) >= s_clamp`. Instead of returning `nan`,
        clamping uses the value of the largest `alpha` that produces a
        non-`nan` coefficient for the given (ai, Re) pair.

        This option is experimental and should be used with caution. Its
        purpose is to mitigate the fictitious, large angles of attack induced
        at the wing tips due to the control points being placed on the lifting
        line. The theory is that if the induced velocity is indeed fictious,
        then the true angle of attack is likely much closer to the standard
        range. By limiting clamping to just the outer `s > s_clamp`, if the
        wing is experiencing a genuinely large angle of attack, then the other
        non-clamped sections will still fail, thus signalling stall conditions.
        If the segments are small the error introduced should be negligible.

    References
    ----------
    .. [1] <NAME>, "Modern Adaptation of Prandtl’s Classic
       Lifting-Line Theory", Journal of Aircraft, 2000

    .. [2] <NAME> Snyder, "A lifting-line approach to estimating
       propeller/wing interactions", 2006

    .. [3] McLeanauth, "Understanding Aerodynamics - Arguing from the Real
       Physics", 2013, p382

    Notes
    -----
    This implementation uses a single distribution for the entire span, which
    is suitable for parafoils, which is a continuous lifting surface, but for
    wings with left and right segments separated by some discontinuity at the
    root you should distribute the points across each semispan independently.
    See [1]_ for a related discussion.

    This method does suffer an issue where induced velocity goes to infinity as
    the segment lengths tend toward zero (as the number of segments increases,
    or for a poorly chosen point distribution). See [2]_, section 8.2.3.
    """

    def __init__(
        self,
        foil: SimpleFoil,
        v_ref_mag,
        alpha_ref: float,
        s_nodes,
        s_clamp: float | None = None,
    ) -> None:
        self.foil = foil
        self.K = len(s_nodes) - 1  # Number of control points
        self.s_nodes = np.asarray(s_nodes)
        self.nodes = self.foil.surface_xyz(self.s_nodes, 0, 0.25, surface="chord")
        self.s_cps = (self.s_nodes[1:] + self.s_nodes[:-1]) / 2
        self.cps = self.foil.surface_xyz(self.s_cps, 0, 0.25, surface="chord")

        # Enable clamped coefficients at some control points
        if s_clamp is not None:
            self.clamped = np.abs(self.s_cps) >= s_clamp
        else:
            self.clamped = np.full(self.K, False)

        # axis0 are nodes, axis1 are control points, axis2 are vectors or norms
        self.R1 = self.cps - self.nodes[:-1, None]
        self.R2 = self.cps - self.nodes[1:, None]
        self.r1 = np.linalg.norm(self.R1, axis=2)  # Magnitudes of R_{i1,j}
        self.r2 = np.linalg.norm(self.R2, axis=2)  # Magnitudes of R_{i2,j}

        # Wing section orientation unit vectors at each control point
        # Note: Phillip's derivation uses back-left-up coordinates (not `frd`)
        u = -self.foil.section_orientation(self.s_cps).T
        self.u_a, self.u_s, self.u_n = u[0].T, u[1].T, u[2].T

        # Define the differential areas as parallelograms by assuming a linear
        # chord variation between nodes.
        self.dl = self.nodes[1:] - self.nodes[:-1]
        node_chords = self.foil.chord_length(self.s_nodes)
        self.c_avg = (node_chords[1:] + node_chords[:-1]) / 2
        self.dA = self.c_avg * np.linalg.norm(cross3(self.u_a, self.dl), axis=1)

        # Precompute the `v` terms that do not depend on `u_inf`, which are the
        # first bracketed term in Hunsaker Eq:6.
        R1, R2, r1, r2 = self.R1, self.R2, self.r1, self.r2  # Shorthand
        self.v_ij = np.zeros((self.K, self.K, 3))  # Extra terms when `i != j`
        for ij in [(i, j) for i in range(self.K) for j in range(self.K)]:
            if ij[0] == ij[1]:  # Skip singularities when `i == j`
                continue
            self.v_ij[ij] = (
                ((r1[ij] + r2[ij]) * cross3(R1[ij], R2[ij]))  # fmt: skip
                / (r1[ij] * r2[ij] * (r1[ij] * r2[ij] + np.dot(R1[ij], R2[ij])))
            )

        # Precompute a reference solution from a (hopefully easy) base case.
        # Sets an initial "solution" (which isn't actually a solution) just to
        # bootstrap the `__call__` method with an initial `Gamma` value.
        alpha_ref = np.deg2rad(alpha_ref)
        v_mag = np.broadcast_to(v_ref_mag, (self.K, 3))
        v_W2f_ref = -v_mag * np.array([np.cos(alpha_ref), 0, np.sin(alpha_ref)])
        self._reference_solution = {
            "ai": 0,
            "v_W2f": v_W2f_ref,
            "Gamma": np.sqrt(1 - self.s_cps**2),  # Naive ellipse
        }
        try:
            _, _, self._reference_solution = self.__call__(0, v_W2f_ref, 1.2)
        except ConvergenceError as e:
            raise RuntimeError("Phillips: failed to initialize base case")

    def _compute_Reynolds(self, v_W2f, rho_air):
        """Compute the Reynolds number at each control point."""
        # FIXME: verify that using the total airspeed (including spanwise flow)
        #        is okay. A few tests show minimal differences, so for now I'm
        #        not wasting time computing the normal and chordwise flows.
        u = np.linalg.norm(v_W2f, axis=-1)  # airspeed [m/s]
        mu = 1.81e-5  # Standard dynamic viscosity of air
        Re = rho_air * u * self.c_avg / mu
        return Re

    def r_CP2LE(self):
        cps = self.cps.view()
        cps.flags.writeable = False
        return cps

    def _induced_velocities(self, u_inf):
        # 2. Compute the "induced velocity" unit vectors
        #  * ref: Phillips, Eq:6
        R1, R2, r1, r2 = self.R1, self.R2, self.r1, self.r2  # Shorthand
        v = self.v_ij.copy()
        v += (
            cross3(u_inf, R2)
            / (r2 * (r2 - np.einsum("k,ijk->ij", u_inf, R2)))[..., None]
        )
        v -= (
            cross3(u_inf, R1)
            / (r1 * (r1 - np.einsum("k,ijk->ij", u_inf, R1)))[..., None]
        )

        return v / (4 * np.pi)  # axes: (inducer, inducee, 3-vector)

    def _local_velocities(self, v_W2f, Gamma, v):
        # Compute the local fluid velocities
        #  * ref: Hunsaker Eq:5
        #  * ref: Phillips Eq:5 (nondimensional version)
        V = v_W2f + np.einsum("j,jik->ik", Gamma, v)

        # Compute the local angle of attack for each section
        #  * ref: Phillips Eq:9 (dimensional) or Eq:12 (dimensionless)
        V_n = np.einsum("ik,ik->i", V, self.u_n)  # Normal-wise
        V_a = np.einsum("ik,ik->i", V, self.u_a)  # Chordwise
        alpha = np.arctan2(V_n, V_a)

        return V, V_n, V_a, alpha

    def _f(self, Gamma, ai, v_W2f, v, Re):
        # Compute the residual error vector
        #  * ref: Hunsaker Eq:8
        #  * ref: Phillips Eq:14
        V, V_n, V_a, alpha = self._local_velocities(v_W2f, Gamma, v)
        W = cross3(V, self.dl)
        W_norm = np.sqrt(np.einsum("ik,ik->i", W, W))
        Cl = self.foil.sections.Cl(
            self.s_cps,
            ai,
            alpha,
            Re,
            clamp=self.clamped,
        )
        # return 2 * Gamma * W_norm - np.einsum("ik,ik,i,i->i", V, V, self.dA, Cl)
        return 2 * Gamma * W_norm - (V_n**2 + V_a**2) * self.dA * Cl

    def _J(self, Gamma, ai, v_W2f, v, Re, verify_J=False):
        # 7. Compute the Jacobian matrix, `J[ij] = d(f_i)/d(Gamma_j)`
        #  * ref: Hunsaker Eq:11
        V, V_n, V_a, alpha = self._local_velocities(v_W2f, Gamma, v)
        W = cross3(V, self.dl)
        W_norm = np.sqrt(np.einsum("ik,ik->i", W, W))
        Cl = self.foil.sections.Cl(
            self.s_cps,
            ai,
            alpha,
            Re,
            clamp=self.clamped,
        )
        Cl_alpha = self.foil.sections.Cl_alpha(
            self.s_cps,
            ai,
            alpha,
            Re,
            clamp=self.clamped,
        )

        J = 2 * np.diag(W_norm)  # Additional terms for i==j
        J2 = 2 * np.einsum("i,ik,i,jik->ij", Gamma, W, 1 / W_norm, cross3(v, self.dl))
        J3 = (
            np.einsum("i,jik,ik->ij", V_a, v, self.u_n)
            - np.einsum("i,jik,ik->ij", V_n, v, self.u_a)  # fmt: skip
        )
        J3 *= (
            (self.dA * Cl_alpha)[:, None]
            * np.einsum("ik,ik->i", V, V)
            / (V_n**2 + V_a**2)
        )
        J4 = 2 * np.einsum("i,i,ik,jik->ij", self.dA, Cl, V, v)
        J += J2 - J3 - J4

        # Compare the analytical gradient to the finite-difference version
        if verify_J:
            J_true = self._J_finite(Gamma, ai, v_W2f, v, Re)
            if not np.allclose(J, J_true):
                print("\n !!! The analytical Jacobian disagrees. Halting. !!!")
                breakpoint()

        return J

    def _J_finite(self, Gamma, ai, v_W2f, v, Re):
        """Compute the Jacobian using a centered finite distance.

        Useful for checking the analytical gradient.

        Examples
        --------
        >>> J1 = self._J(Gamma, v_W2f, v, ai)
        >>> J2 = self._J_finite(Gamma, v_W2f, v, ai)
        >>> np.allclose(J1, J2)
        True
        """
        # This uses the same method as `scipy.optimize.approx_fprime`, but that
        # function only works for scalar-valued functions.
        JT = np.empty((self.K, self.K))  # Jacobian transpose  (J_ji)
        eps = np.sqrt(np.finfo(float).eps)

        # Build the Jacobian column-wise (row-wise of the tranpose)
        Gp, Gm = Gamma.copy(), Gamma.copy()
        for k in range(self.K):
            Gp[k], Gm[k] = Gamma[k] + eps, Gamma[k] - eps
            fp = self._f(Gp, ai, v_W2f, v, Re)
            fm = self._f(Gm, ai, v_W2f, v, Re)
            JT[k] = (fp - fm) / (2 * eps)
            Gp[k], Gm[k] = Gamma[k], Gamma[k]

        return JT.T

    def _solve_circulation(self, ai, v_W2f, Re, Gamma0):
        """
        Solve for the spanwise circulation distribution.

        Parameters
        ----------
        ai : array of float, shape (K,) [radians]
            Airfoil indices.
        v_W2f : array of float, shape (K,) [m/s]
            Relative wind velocity at each control point.
        Re : array of float, shape (K,)
            Reynolds number at each segment
        Gamma0 : array of float, shape (K,)
            The initial proposal

        Returns
        -------
        Gamma : array of float, shape (K,)
            Circulation strengths of each segment.
        v : array, shape (K,K,3) [m/s]
            Induced velocities between each segment, indexed as (inducer,
            inducee).
        """
        v_mid = v_W2f[self.K // 2]
        u_inf = v_mid / np.linalg.norm(v_mid)  # FIXME: what if PQR != 0?
        v = self._induced_velocities(u_inf)
        args = (ai, v_W2f, v, Re)
        res = scipy.optimize.root(self._f, Gamma0, args, jac=self._J, tol=1e-4)

        if not res["success"]:
            raise ConvergenceError

        return res["x"], v

    def __call__(
        self,
        ai,
        v_W2f,
        rho_air,
        *,
        reference_solution: dict | None = None,
        max_splits: int = 10,
    ):
        v_W2f = np.broadcast_to(v_W2f, (self.K, 3))
        Re = self._compute_Reynolds(v_W2f, rho_air)

        if reference_solution is None:
            reference_solution = self._reference_solution

        ai_ref = reference_solution["ai"]
        v_W2f_ref = reference_solution["v_W2f"]
        Gamma_ref = reference_solution["Gamma"]

        # Try to solve for the target (`Gamma` as a function of `v_W2f` and
        # `ai`) directly using the `reference_solution`. If that fails, pick a
        # point between the target and the reference, solve for that easier
        # case, then use its solution as the new starting point for the next
        # target. Repeat for intermediate targets until either solving for the
        # original target, or exceeding `max_splits`.
        target_backlog = []  # Stack of pending targets
        num_splits = 0
        while True:
            try:
                Gamma, v = self._solve_circulation(ai, v_W2f, Re, Gamma_ref)
            except ConvergenceError:
                if num_splits == max_splits:
                    raise ConvergenceError("max splits reached")
                num_splits += 1
                target_backlog.append((ai, v_W2f))
                P = 0.5  # Ratio, a point between the reference and the target
                ai = (1 - P) * ai_ref + P * ai
                v_W2f = (1 - P) * v_W2f_ref + P * v_W2f
                continue

            ai_ref = ai
            v_W2f_ref = v_W2f
            Gamma_ref = Gamma

            if target_backlog:
                ai, v_W2f = target_backlog.pop()
            else:
                break

        V, V_n, V_a, alpha = self._local_velocities(v_W2f, Gamma, v)

        # Compute the inviscid forces using the 3D vortex lifting law
        #  * ref: Hunsaker Eq:1
        #  * ref: Phillips Eq:4
        dF_inviscid = Gamma * cross3(V, self.dl).T

        # Compute the viscous forces.
        #  * ref: Hunsaker Eq:17
        #
        # The equation in the paper uses the "characteristic chord", but I
        # believe that is a mistake; it produces *massive* drag. Here I use the
        # section area like they do in "MachUp_Py" (see where they compute
        # `f_parasite_mag` in `llmodel.py:LLModel:_compute_forces`).
        Cd = self.foil.sections.Cd(
            self.s_cps,
            ai,
            alpha,
            Re,
            clamp=self.clamped,
        )
        V2 = np.einsum("ik,ik->i", V, V)
        u_drag = V.T / np.sqrt(V2)
        dF_viscous = 0.5 * V2 * self.dA * Cd * u_drag

        # The total forces applied at each control point
        dF = dF_inviscid + dF_viscous

        # Compute the section moments.
        #  * ref: Hunsaker Eq:19
        #  * ref: Phillips Eq:28
        #
        # These are strictly the section moments caused by airflow around the
        # section. It does not include moments about the aircraft reference
        # point (commonly the center of gravity); those extra moments must be
        # calculated by the wing.
        #  * ref: Hunsaker Eq:19
        #  * ref: Phillips Eq:28
        Cm = self.foil.sections.Cm(
            self.s_cps,
            ai,
            alpha,
            Re,
            clamp=self.clamped,
        )
        dM = -0.5 * V2 * self.dA * self.c_avg * Cm * self.u_s.T

        solution = {
            "ai": ai,
            "v_W2f": v_W2f_ref,
            "Gamma": Gamma_ref,
        }

        # print("\nFinished `Phillips.__call__`")
        # breakpoint()

        dF *= rho_air
        dM *= rho_air

        return dF.T, dM.T, solution


<reponame>smujiang/WSIPenMarkingRemoval<gh_stars>1-10
#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Compare two or more images using MSE, PSNR, SNR, SSIM, UQI, PBVIF, MSSIM,
NQM and WSNR metrics.

For usage and a list of options, try this:
$ ./pymetrikz -h

This program and its regression test suite live here:
http://www.sawp.com.br/projects/pymetrikz"""

import numpy as __n
from scipy.ndimage.filters import gaussian_filter as __gaussian_filter
from scipy.ndimage.filters import convolve as __convolve
from scipy.ndimage.filters import correlate as __correlate
from scipy.fftpack import fftshift as __fftshift

__author__ = "<NAME> <<EMAIL>>"
__copyright__ = "Copyright (c) 2011-2014 <NAME>"
__license__ = "GPLv2"

def mse(reference, query):
    """Computes the Mean Square Error (MSE) of two images.

    value = mse(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : MSE value
    """
    (ref, que) = (reference.astype('double'), query.astype('double'))
    diff = ref - que
    square = (diff ** 2)
    mean = square.mean()
    return mean

def rmse(reference, query):
    msev = mse(reference, query)
    return __n.sqrt(msev)

def psnr(reference, query, normal=255):
    """Computes the Peak Signal-to-Noise-Ratio (PSNR).

    value = psnr(reference, query, normalization=255)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.
    normal   : normalization value (255 for 8-bit image

    Return
    ----------
    value    : PSNR value
    """
    normalization = float(normal)
    msev = mse(reference, query)
    if msev != 0:
        value = 10.0 * __n.log10(normalization * normalization / msev)
    else:
        value = float("inf")
    return value

def snr(reference, query):
    """Computes the Signal-to-Noise-Ratio (SNR).

    value = snr(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : SNR value
    """
    signal_value = (reference.astype('double') ** 2).mean()
    msev = mse(reference, query)
    if msev != 0:
        value = 10.0 * __n.log10(signal_value / msev)
    else:
        value = float("inf")
    return value

def ssim(reference, query):
    """Computes the Structural SIMilarity Index (SSIM).

    value = ssim(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : SSIM value
    """
    def __get_kernels():
        k1, k2, l = (0.01, 0.03, 255.0)
        kern1, kern2 = map(lambda x: (x * l) ** 2, (k1, k2))
        return kern1, kern2

    def __get_mus(i1, i2):
        mu1, mu2 = map(lambda x: __gaussian_filter(x, 1.5), (i1, i2))
        m1m1, m2m2, m1m2 = (mu1 * mu1, mu2 * mu2, mu1 * mu2)
        return m1m1, m2m2, m1m2

    def __get_sigmas(i1, i2, delta1, delta2, delta12):
        f1 = __gaussian_filter(i1 * i1, 1.5) - delta1
        f2 = __gaussian_filter(i2 * i2, 1.5) - delta2
        f12 = __gaussian_filter(i1 * i2, 1.5) - delta12
        return f1, f2, f12

    def __get_positive_ssimap(C1, C2, m1m2, mu11, mu22, s12, s1s1, s2s2):
        num = (2 * m1m2 + C1) * (2 * s12 + C2)
        den = (mu11 + mu22 + C1) * (s1s1 + s2s2 + C2)
        return num / den

    def __get_negative_ssimap(C1, C2, m1m2, m11, m22, s12, s1s1, s2s2):
        (num1, num2) = (2.0 * m1m2 + C1, 2.0 * s12 + C2)
        (den1, den2) = (m11 + m22 + C1, s1s1 + s2s2 + C2)
        ssim_map = __n.ones(img1.shape)
        indx = (den1 * den2 > 0)
        ssim_map[indx] = (num1[indx] * num2[indx]) / (den1[indx] * den2[indx])
        indx = __n.bitwise_and(den1 != 0, den2 == 0)
        ssim_map[indx] = num1[indx] / den1[indx]
        return ssim_map

    (img1, img2) = (reference.astype('double'), query.astype('double'))
    (m1m1, m2m2, m1m2) = __get_mus(img1, img2)
    (s1, s2, s12) = __get_sigmas(img1, img2, m1m1, m2m2, m1m2)
    (C1, C2) = __get_kernels()
    if C1 > 0 and C2 > 0:
        ssim_map = __get_positive_ssimap(C1, C2, m1m2, m1m1, m2m2, s12, s1, s2)
    else:
        ssim_map = __get_negative_ssimap(C1, C2, m1m2, m1m1, m2m2, s12, s1, s2)
    ssim_value = ssim_map.mean()
    return ssim_value

def uqi(reference, query):
    """Computes the Universal Quality Index (UQI).

    value = uqi(reference, query

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : UQI value
    """
    def __conv(x):
        window = __n.ones((BLOCK_SIZE, BLOCK_SIZE))
        if len(x.shape) < 3:
            return __convolve(x, window)
        else:
            channels = x.shape[2]
            f = [__convolve(x[:, :, c], window) for c in range(channels)]
            return __n.array(f)

    def __get_filtered(im1, im2, BLOCK_SIZE):
        (im1im1, im2im2, im1im2) = (im1 * im1, im2 * im2, im1 * im2)
        (b1, b2, b3, b4, b5) = map(__conv, (im1, im2, im1im1, im2im2, im1im2))
        (b6, b7) = (b1 * b2, b1 * b1 + b2 * b2)
        return (b1, b2, b3, b4, b5, b6, b7)

    def __get_quality_map(b1, b2, b3, b4, b5, b6, b7, BLOCK_SIZE):
        N = BLOCK_SIZE * BLOCK_SIZE
        numerator = 4.0 * (N * b5 - b6) * b6
        denominator1 = N * (b3 + b4) - b7
        denominator = denominator1 * b7
        index = __n.bitwise_and(denominator1 == 0, b7 != 0)
        quality_map = __n.ones(denominator.shape)
        quality_map[index] = 2.0 * b6[index] / b7[index]
        index = (denominator != 0)
        quality_map[index] = numerator[index] / denominator[index]
        return quality_map[index]

    BLOCK_SIZE = 8
    (img1, img2) = (reference.astype('double'), query.astype('double'))
    (b1, b2, b3, b4, b5, b6, b7) = __get_filtered(img1, img2, BLOCK_SIZE)
    quality_map = __get_quality_map(b1, b2, b3, b4, b5, b6, b7, BLOCK_SIZE)
    value = quality_map.mean()
    return value

def pbvif(reference, query):
    """Computes the Pixel-Based Visual Information Fidelity (PB-VIF).

    value = pbvif(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : PB-VIF value
    """
    def __get_sigma(win, ref, dist, mu1_sq, mu2_sq, mu1_mu2):
        sigma1_sq = __filter2(win, ref * ref) - mu1_sq
        sigma2_sq = __filter2(win, dist * dist) - mu2_sq
        sigma12 = __filter2(win, ref * dist) - mu1_mu2
        (sigma1_sq[sigma1_sq < 0], sigma2_sq[sigma2_sq < 0]) = (0.0, 0.0)
        return (sigma2_sq, sigma12, sigma1_sq)

    def __get_normalized(s1s1, s2s2, s1s2):
        g = s1s2 / (s1s1 + 1e-10)
        sv_sq = s2s2 - g * s1s2
        g[s1s1 < 1e-10] = 0
        sv_sq[s1s1 < 1e-10] = s2s2[s1s1 < 1e-10]
        s1s1[s1s1 < 1e-10] = 0
        g[s2s2 < 1e-10] = 0
        sv_sq[s2s2 < 1e-10] = 0
        sv_sq[g < 0] = s2s2[g < 0]
        g[g < 0] = 0
        sv_sq[sv_sq <= 1e-10] = 1e-10
        return (g, sv_sq)

    def __get_num(s1s1, sv_sq, sigma_nsq, g):
        normg = (g ** 2) * s1s1 / (sv_sq + sigma_nsq)
        snr = __n.log10(1.0 + normg).sum()
        return snr

    def __get_den(s1s1, sigma_nsq):
        snr = __n.log10(1.0 + s1s1 / sigma_nsq)
        return snr.sum()

    def __get_num_den_level(ref, dist, scale):
        sig = 2.0
        N = (2.0 ** (4 - scale + 1.0)) + 1.0
        win = __get_gaussian_kernel(N, N / 5.0)
        if scale > 1:
            ref = __filter2(win, ref)
            dist = __filter2(win, dist)
            ref = ref[::2, ::2]
            dist = dist[::2, ::2]
        (mu1, mu2) = (__filter2(win, ref), __filter2(win, dist))
        (m1m1, m2m2, m1m2) = (mu1 * mu1, mu2 * mu2, mu1 * mu2)
        (s2s2, s1s2, s1s1) = __get_sigma(win, ref, dist, m1m1, m2m2, m1m2)
        (g, svsv) = __get_normalized(s1s1, s2s2, s1s2)
        (num, den) = (__get_num(s1s1, svsv, sig, g), __get_den(s1s1, sig))
        return (num, den)

    (ref, dist) = (reference.astype('double'), query.astype('double'))
    zipped = map(lambda x: __get_num_den_level(ref, dist, x), range(1, 5))
    (nums, dens) = zip(*zipped)
    value = sum(nums) / sum(dens)
    return value

def mssim(reference, query):
    """Computes the Multi-Scale SSIM Index (MSSIM).

    value = mssim(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : MSSIM value
    """
    def __get_filt_kern():
        n = [131, -199, -101, 962, 932, 962, -101, -199, 131]
        d = [3463, 8344, 913, 2549, 1093, 2549, 913, 8344, 3463]
        num = __n.matrix(n).T
        den = __n.matrix(d).T
        lod = num.astype('double') / den.astype('double')
        lpf = __n.dot(lod, lod.T)
        return lpf / lpf.sum()

    def __get_ssim(img1, img2, K):
        comp_ssim = __ssim_modified(img1, img2, K)[1]
        return (comp_ssim[1], comp_ssim[2])

    def __get_MVR(img1, img2, K, nlevs):
        (ssim_v, ssim_r) = (__n.zeros((nlevs, 1)), __n.zeros((nlevs, 1)))
        (ssim_v[0], ssim_r[0]) = __get_ssim(img1, img2, K)
        filt_kern = __get_filt_kern()
        for s in range(nlevs - 1):
            (img1, img2) = map(lambda x: __filter2(filt_kern, x), (img1, img2))
            (img1, img2) = (img1[::2, ::2], img2[::2, ::2])
            comp_ssim = __ssim_modified(img1, img2, K)[1]
            ssim_m = comp_ssim[0]
            ssim_v[s + 1] = comp_ssim[1]
            ssim_r[s + 1] = comp_ssim[2]
        return (ssim_m, ssim_v, ssim_r)

    def __calc_mssim_mvr(img1, img2):
        (K, weights) = ((0.01, 0.03), (0.0448, 0.2856, 0.3001, 0.2363, 0.1333))
        (alpha, beta, lvl) = (0.1333, __n.matrix(weights).T, len(weights))
        (ssim_m, ssim_v, ssim_r) = __get_MVR(img1, img2, K, lvl)
        m = ssim_m ** alpha
        v = (ssim_v ** beta).prod()
        r = (ssim_r ** beta).prod()
        return (m, v, r)

    (ref, quer) = (reference.astype('double'), query.astype('double'))
    ssim_mvr = __n.matrix(__calc_mssim_mvr(ref, quer))
    value = ssim_mvr.prod()
    return value

def __filter2(B, X, shape='nearest'):
    B2 = __n.rot90(__n.rot90(B))
    if len(X.shape) < 3:
        return __correlate(X, B2, mode=shape)
    else:
        channels = X.shape[2]
        f = [__correlate(X[:, :, c], B2, mode=shape) for c in range(channels)]
        return __n.array(f)

def __get_gaussian_kernel(N=15, sigma=1.5):
    (H, W) = ((N - 1) / 2, (N - 1) / 2)
    std = sigma
    (y, x) = __n.mgrid[-H:H + 1, -W:W + 1]
    arg = -(x * x + y * y) / (2.0 * std * std)
    h = __n.exp(arg)
    index = h < __n.finfo(float).eps * h.max(0)
    h[index] = 0
    sumh = h.sum()
    if sumh != 0:
        h = h / sumh
    return h

def __ssim_modified(reference, query, K):
    def __get_kern(K):
        L = 255
        kern = map(lambda x: (x * L) ** 2, K)
        return (kern[0], kern[1])

    def __get_filtering_window():
        window = __get_gaussian_kernel(11, 1.5)
        return window / window.sum()

    def __get_mus(img1, img2, window):
        (mu1, mu2) = map(lambda x: __filter2(window, x), (img1, img2))
        (m1m1, m2m2, m1m2) = (mu1 * mu1, mu2 * mu2, mu1 * mu2)
        return (mu1, mu2, m1m1, m2m2, m1m2)

    def __get_sigmas(img1, img2, window, m1m1, m2m2, m1m2):
        s1s1 = __filter2(window, img1 * img1) - m1m1
        s2s2 = __filter2(window, img2 * img2) - m2m2
        s12 = __filter2(window, img1 * img2) - m1m2
        (s1, s2) = map(__n.sqrt, (__n.abs(s1s1), __n.abs(s2s2)))
        return (s1s1, s2s2, s1, s2, s12)

    def __MVR_pos_kern(m, kern, s, s_square):
        (m11, m22, m12) = m
        (k1, k2) = kern
        (s1, s2) = s
        (s1s1, s2s2, s12) = s_square
        M = (2.0 * m12 + k1) / (m11 + m22 + k1)
        V = (2.0 * s1 * s2 + k2) / (s1s1 + s2s2 + k2)
        R = (s12 + k2 / 2.0) / (s1 * s2 + k2 / 2.0)
        return (M, V, R)

    def __MVR_neg_kern(m, s, s_square):
        def __calcM(mu1, m11, m22, m12):
            ssim_ln = 2.0 * m12
            ssim_ld = m11 + m22
            index_l = ssim_ld > 0
            M = __n.ones(mu1.shape)
            M[index_l] = ssim_ln[index_l] / ssim_ld[index_l]
            return M

        def __calcV(mu1, s1, s2, s11, s22):
            ssim_cn = 2.0 * s1 * s2
            ssim_cd = s11 + s22
            V = __n.ones(mu1.shape)
            index_c = ssim_cd > 0
            V[index_c] = ssim_cn[index_c] / ssim_cd[index_c]
            return V

        def __calcR(mu1, s1, s2, s12):
            (ssim_sn, ssim_sd) = (s12, s1 * s2)
            R = __n.ones(mu1.shape)
            (index1, index2) = (s1 > 0, s2 > 0)
            index_s1 = index1 * index2 > 0
            R[index_s1] = ssim_sn[index_s1] / ssim_sd[index_s1]
            index_s2 = index1 * __n.logical_not(index2) > 0
            R[index_s2] = 0.0
            return R

        (mu1, mu2, m11, m22, m12) = m
        (s1, s2) = s
        (s11, s22, s12) = s_square
        M = __calcM(mu1, m11, m22, m12)
        V = __calcV(mu1, s1, s2, s11, s22)
        R = __calcR(mu1, s1, s2, s12)
        return (M, V, R)

    def __get_composition_vector(img1, img2):
        filt = __get_filtering_window()
        (mu1, mu2, m11, m22, m12) = __get_mus(img1, img2, filt)
        (s11, s22, s1, s2, s12) = __get_sigmas(img1, img2, filt, m11, m22, m12)
        (kern1, kern2) = __get_kern(K)
        if kern1 > 0 and kern2 > 0:
            (m, kern, s) = ((m11, m22, m12), (kern1, kern2), (s1, s2))
            s_square = (s11, s22, s12)
            (M, V, R) = __MVR_pos_kern(m, kern, s, s_square)
        else:
            (m, s) = ((mu1, mu2, m11, m22, m12), (s1, s2))
            s_square = (s11, s22, s12)
            (M, V, R) = __MVR_neg_kern(m, s, s_square)
        return (M, V, R)

    def __get_ssim_map(M, V, R):
        ssim_map = M * V * R
        return ssim_map

    def __get_ssim_from_composition_vector(M, V, R):
        ssim_map = __get_ssim_map(M, V, R)
        ssim = ssim_map.mean()
        return ssim

    (img1, img2) = reference.astype('double'), query.astype('double')
    (M, V, R) = __get_composition_vector(img1, img2)
    composite_mean_vector = (M.mean(), V.mean(), R.mean())
    ssim = __get_ssim_from_composition_vector(M, V, R)
    return (ssim, composite_mean_vector)

def __convert_to_luminance(x):
    return __n.dot(x[..., :3], [0.299, 0.587, 0.144]).astype('double')

def nqm(reference, query):
    """Computes the NQM metric.

    value = nqm(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : NQM value
    """
    def __ctf(f_r):
        """ Bandpass Contrast Threshold Function for RGB"""
        (gamma, alpha) = (0.0192 + 0.114 * f_r, (0.114 * f_r) ** 1.1)
        beta = __n.exp(-alpha)
        num = 520.0 * gamma * beta
        return 1.0 / num

    def _get_masked(c, ci, a, ai, i):
        (H, W) = c.shape
        (c, ci, ct) = (c.flatten(1), ci.flatten(1), __ctf(i))
        ci[abs(ci) > 1.0] = 1.0
        T = ct * (0.86 * ((c / ct) - 1.0) + 0.3)
        (ai, a, a1) = (ai.flatten(1), a.flatten(1), (abs(ci - c) - T) < 0.0)
        ai[a1] = a[a1]
        return ai.reshape(H, W)

    def __get_thresh(x, T, z, trans=True):
        (H, W) = x.shape
        if trans:
            (x, z) = (x.flatten(1).T, z.flatten())
        else:
            (x, z) = (x.flatten(1), z.flatten(1))
        z[abs(x) < T] = 0.0
        return z.reshape(H, W)

    def __decompose_cos_log_filter(w1, w2, phase=__n.pi):
        return 0.5 * (1 + __n.cos(__n.pi * __n.log2(w1 + w2) - phase))

    def __get_w(r):
        w = [(r + 2) * ((r + 2 <= 4) * (r + 2 >= 1))]
        w += [r * ((r <= 4) * (r >= 1))]
        w += [r * ((r >= 2) * (r <= 8))]
        w += [r * ((r >= 4) * (r <= 16))]
        w += [r * ((r >= 8) * (r <= 32))]
        w += [r * ((r >= 16) * (r <= 64))]
        return w

    def __get_u(r):
        u = [4 * (__n.logical_not((r + 2 <= 4) * (r + 2 >= 1)))]
        u += [4 * (__n.logical_not((r <= 4) * (r >= 1)))]
        u += [0.5 * (__n.logical_not((r >= 2) * (r <= 8)))]
        u += [4 * (__n.logical_not((r >= 4) * (r <= 16)))]
        u += [0.5 * (__n.logical_not((r >= 8) * (r <= 32)))]
        u += [4 * (__n.logical_not((r >= 16) * (r <= 64)))]
        return u

    def __get_G(r):
        (w, u) = (__get_w(r), __get_u(r))
        phase = [__n.pi, __n.pi, 0.0, __n.pi, 0.0, __n.pi]
        dclf = __decompose_cos_log_filter
        return [dclf(w[i], u[i], phase[i]) for i in range(len(phase))]

    def __compute_fft_plane_shifted(ref, query):
        (x, y) = ref.shape
        (xplane, yplane) = __n.mgrid[-y / 2:y / 2, -x / 2:x / 2]
        plane = (xplane + 1.0j * yplane)
        r = abs(plane)
        G = __get_G(r)
        Gshifted = map(__fftshift, G)
        return [Gs.T for Gs in Gshifted]

    def __get_c(a, l_0):
        c = [a[0] / l_0]
        c += [a[1] / (l_0 + a[0])]
        c += [a[2] / (l_0 + a[0] + a[1])]
        c += [a[3] / (l_0 + a[0] + a[1] + a[2])]
        c += [a[4] / (l_0 + a[0] + a[1] + a[2] + a[3])]
        return c

    def __get_ci(ai, li_0):
        ci = [ai[0] / (li_0)]
        ci += [ai[1] / (li_0 + ai[0])]
        ci += [ai[2] / (li_0 + ai[0] + ai[1])]
        ci += [ai[3] / (li_0 + ai[0] + ai[1] + ai[2])]
        ci += [ai[4] / (li_0 + ai[0] + ai[1] + ai[2] + ai[3])]
        return ci

    def __compute_contrast_images(a, ai, l, li):
        ci = __get_ci(ai, li)
        c = __get_c(a, l)
        return (c, ci)

    def __get_detection_thresholds():
        viewing_angle = (1.0 / 3.5) * (180.0 / __n.pi)
        rotations = [2.0, 4.0, 8.0, 16.0, 32.0]
        return map(lambda x: __ctf(x / viewing_angle), rotations)

    def __get_account_for_supra_threshold_effects(c, ci, a, ai):
        r = range(len(a))
        return [_get_masked(c[i], ci[i], a[i], ai[i], i + 1) for i in r]

    def __apply_detection_thresholds(c, ci, d, a, ai):
        A = [__get_thresh(c[i], d[i], a[i], False) for i in range(len(a))]
        AI = [__get_thresh(ci[i], d[i], ai[i], True) for i in range(len(a))]
        return (A, AI)

    def __reconstruct_images(A, AI):
        return map(lambda x: __n.add.reduce(x), (A, AI))

    def __compute_quality(imref, imquery):
        return snr(imref, imquery)

    def __get_ref_basis(ref_fft, query_fft, GS):
        (L_0, LI_0) = map(lambda x: GS[0] * x, (ref_fft, query_fft))
        (l_0, li_0) = map(lambda x: __n.real(__n.fft.ifft2(x)), (L_0, LI_0))
        return (l_0, li_0)

    def __compute_inverse_convolution(convolved_fft, GS):
        convolved = [GS[i] * convolved_fft for i in range(1, len(GS))]
        return map(lambda x: __n.real(__n.fft.ifft2(x)), convolved)

    def __correlate_in_fourier_domain(ref, query):
        (ref_fft, query_fft) = map(lambda x: __n.fft.fft2(x), (ref, query))
        GS = __compute_fft_plane_shifted(ref, query)
        (l_0, li_0) = __get_ref_basis(ref_fft, query_fft, GS)
        a = __compute_inverse_convolution(ref_fft, GS)
        ai = __compute_inverse_convolution(query_fft, GS)
        return (a, ai, l_0, li_0)

    def __get_correlated_images(ref, query):
        (a, ai, l_0, li_0) = __correlate_in_fourier_domain(ref, query)
        (c, ci) = __compute_contrast_images(a, ai, l_0, li_0)
        d = __get_detection_thresholds()
        ai = __get_account_for_supra_threshold_effects(c, ci, a, ai)
        return __apply_detection_thresholds(c, ci, d, a, ai)

    if not len(reference.shape) < 3:
        reference = __convert_to_luminance(reference)
        query = __convert_to_luminance(query)
    (A, AI) = __get_correlated_images(reference, query)
    (y1, y2) = __reconstruct_images(A, AI)
    y = __compute_quality(y1, y2)
    return y

def wsnr(reference, query):
    """Computes the Weighted Signal to Noise Ratio (WSNR) metric.

    value = wsnr(reference, query)

    Parameters
    ----------
    reference: original image data.
    query    : modified image data to be compared.

    Return
    ----------
    value    : wsnr value
    """
    def __genetate_meshgrid(x, y):
        f = lambda u: u / 2 + 0.5 - 1
        (H, W) = map(f, (x, y))
        return (H, W)

    def __create_complex_planes(x, y):
        (H, W) = __genetate_meshgrid(x, y)
        (xplane, yplane) = __n.mgrid[-H:H + 1, -W:W + 1]
        return (xplane, yplane)

    def __get_evaluated_contrast_sensivity(plane):
        w = 0.7
        angle = __n.angle(plane)
        return ((1.0 - w) / 2.0) * __n.cos(4.0 * angle) + (1.0 + w) / 2.0

    def __get_radial_frequency(x, y):
        (xplane, yplane) = __create_complex_planes(x, y)
        nfreq = 60
        plane = (xplane + 1.0j * yplane) / x * 2.0 * nfreq
        s = __get_evaluated_contrast_sensivity(plane)
        radfreq = abs(plane) / s
        return radfreq

    def __generate_CSF(radfreq):
        a = -((0.114 * radfreq) ** 1.1)
        csf = 2.6 * (0.0192 + 0.114 * radfreq) * __n.exp(a)
        f = radfreq < 7.8909
        csf[f] = 0.9809
        return csf

    def __weighted_fft_domain(ref, quer, csf):
        err = ref.astype('double') - quer.astype('double')
        err_wt = __fftshift(__n.fft.fft2(err)) * csf
        im = __n.fft.fft2(ref)
        return (err, err_wt, im)

    def __get_weighted_error_power(err_wt):
        return (err_wt * __n.conj(err_wt)).sum()

    def __get_signal_power(im):
        return (im * __n.conj(im)).sum()

    def __get_ratio(mss, mse):
        if mse != 0:
            ratio = 10.0 * __n.log10(mss / mse)
        else:
            ratio = float("inf")
        return __n.real(ratio)

    if not len(reference.shape) < 3:
        reference = __convert_to_luminance(reference)
        query = __convert_to_luminance(query)
    size = reference.shape
    (x, y) = (size[0], size[1])
    radfreq = __get_radial_frequency(x, y)
    csf = __generate_CSF(radfreq)
    (err, err_wt, im) = __weighted_fft_domain(reference, query, csf)
    mse = __get_weighted_error_power(err_wt)
    mss = __get_signal_power(im)
    ratio = __get_ratio(mss, mse)
    return ratio