import matplotlib.pyplot as plt
import numpy as np

import pandas as pd
import xlsxwriter


cuts_all = {
    'Spanish': {'2026': {'C': [49, 50], 'B': [73, 63], 'T': [99, 80], 'L': [0, 20], 'U': [100, 80]}},
    'French':  {'2010': {'C': [40, 50], 'B': [68, 59], 'T': [99, 80], 'L': [0, 20], 'U': [100, 80]}},
    'German':  {'2010': {'C': [34, 50], 'B': [56, 60], 'T': [86, 80], 'L': [0, 20], 'U': [87, 80]}}
}

col_dict = {
    'raw': 'Raw Score', 
    'ss_lin':'SS Linear',
    'ss_elb':'SS Segments',
    'ss_lin_ur':'SS Linear (unrounded)',
    'ss_elb_ur':'SS Segments (unrounded)',
}

def ab(top, bot, cuts):
    _a = (cuts[top][1] - cuts[bot][1]) / (cuts[top][0] - cuts[bot][0])
    _b = cuts[top][1] - _a * cuts[top][0]

    return _a, _b


def round_clip(a, low, high):
    return np.clip(np.round(a, 0), low, high)


def conversion(cuts):
    a, b = ab('T', 'C', cuts)
    a1, b1 = ab('B', 'C', cuts)
    a2, b2 = ab('T', 'B', cuts)
    raw = np.arange(cuts['L'][0], cuts['U'][0]+1) 
    raw1 = np.arange(cuts['L'][0], cuts['B'][0])
    raw2 = np.arange(cuts['B'][0], cuts['U'][0]+1)
    ss_lin_ur = raw * a + b
    ss_elb_ur = np.append(raw1 * a1 + b1, raw2 * a2 + b2)
    data = pd.DataFrame({
        'raw': raw,
        'ss_lin': round_clip(ss_lin_ur, cuts['L'][1], cuts['U'][1]),
        'ss_elb': round_clip(ss_elb_ur, cuts['L'][1], cuts['U'][1]),
        'ss_lin_ur': ss_lin_ur,
        'ss_elb_ur': ss_elb_ur
    })
    data = data.sort_values(by="raw", ascending=False)

    return data   

    
def plot(data, subject, year, cuts):
    fig, ax = plt.subplots()
    ax.plot(data.raw, data.ss_lin_ur, label='Linear')
    ax.plot(data.raw, data.ss_elb_ur, color='magenta', label='Segments')
    ax.plot([cuts['L'][0], cuts['U'][0]], [cuts['C'][1], cuts['C'][1]], color='green', linewidth=0.5) 
    ax.plot([cuts['L'][0], cuts['U'][0]], [cuts['B'][1], cuts['B'][1]], color='orange', linewidth=0.5)
    ax.set_xlabel('Raw Score')
    ax.set_ylabel('Scale Score')
    ax.set_title(f'Conversions: CLEP {subject} ({year})')
    ax.legend()
    plt.tight_layout()
    plt.savefig(f'output/CLEP_{subject}_conversion_{year}.pdf')


def save_excel(data, subject, year, cuts):
    workbook = xlsxwriter.Workbook(f'output/CLEP_{subject}_conversion_{year}.xlsx')
    worksheet = workbook.add_worksheet('conversion')
    bold = workbook.add_format({'bold': True})
    cut_red0 = workbook.add_format({'bold': True, 'font_color': 'red'})
    cut_red2 = workbook.add_format({'bold': True, 'font_color': 'red', 'num_format': '#.00'})
    dec0 = workbook.add_format({'num_format': '#0'})
    dec2 = workbook.add_format({'num_format': '#.0'})
    i, j = 0, 0
    for col in data.columns:
        for row in data[col]:
            fmt = dec0 
            val = data[col].iloc[j]
            raw = data.raw.iloc[j]
            if j == 0: # column header
                worksheet.write(j, i, col_dict[col], bold)
            if (raw == cuts['C'][0] or raw == cuts['B'][0]):
                if i<3:
                    fmt = cut_red0
                else:
                    fmt = cut_red2
            else:
                if i>2:
                    fmt = dec2
            worksheet.write(j+1, i, val, fmt)
            j += 1
        i += 1
        j = 0
    workbook.close()


def main(subject, year):
    cuts = cuts_all[subject][year]
    data = conversion(cuts)
    plot(data, subject, year, cuts)
    save_excel(data, subject, year, cuts)


if __name__ == "__main__":
    main('Spanish', '2026')