const startInput = document.getElementById('startDate');
const endInput = document.getElementById('endDate');
const cutOmicronInput = document.getElementById('cutOmicron');

const hashParams = new URLSearchParams(window.location.hash.substring(1));
startInput.value = hashParams.get('start');
endInput.value = hashParams.get('end');
cutOmicronInput.checked = hashParams.get('cutOmicron') === 'true';

const OMICRON_START_DATE = '2021-11-01';
const OMICRON_END_DATE = '2022-02-31';

fetch('/data.json')
    .then(resp => resp.json())
    .then(plot)
    .catch(err => {
        console.log(err);
    });

function average(window) {
    let y_sum = 0;
    let count = 0;
    for (const sample of window) {
        if (sample['y'] !== undefined) {
            y_sum += sample['y'];
            count += 1;
        }
    }
    if (count === 0) {
        return undefined;
    }
    return y_sum / count;
}

function localLinearRegression(slopeWindowSize, averageWindowSize) {
    return (window) => {
        let x_sum = 0;
        let y_sum = 0;
        let count = 0;
        for (const sample of window.slice(-slopeWindowSize)) {
            if (sample['y'] !== undefined) {
                x_sum += sample['x'];
                y_sum += sample['y'];
                count += 1;
            }
        }
        if (count === 0) {
            return undefined;
        }
        const x_avg = x_sum / count;
        const y_avg = y_sum / count;

        let numerator = 0;
        let denominator = 0;
        for (const sample of window) {
            if (sample['y'] !== undefined) {
                const x_err = sample['x'] - x_avg;
                const y_err = sample['y'] - y_avg;
                numerator += x_err * y_err;
                denominator += x_err * x_err;
            }
        }
        if (denominator === 0) {
            return window.at(-1)['y'];
        }
        const slope = numerator / denominator;

        let short_x_sum = 0;
        let short_y_sum = 0;
        let short_count = 0;
        for (const sample of window.slice(-averageWindowSize)) {
            if (sample['y'] !== undefined) {
                short_x_sum += sample['x'];
                short_y_sum += sample['y'];
                short_count += 1;
            }
        }
        if (short_count === 0) {
            return y_avg + (window.at(-1)['x'] - x_avg) * slope;
        }
        const short_x_avg = short_x_sum / short_count;
        const short_y_avg = short_y_sum / short_count;
        return short_y_avg + (window.at(-1)['x'] - short_x_avg) * slope;
    };
}

function smooth(data, column, windowSize, smoothFunc) {
    const smoothed = [];
    const window = [];
    for (const row of data) {
        window.push({ x: new Date(row['x']).getTime(), y: row[column] });
        if (window.length > windowSize) {
            window.shift();
        }
        smoothed.push({ x: row['x'], y: smoothFunc(window) });
    }
    return smoothed;
}

function extractWithErrorBars(data, region) {
    const extracted = [];
    for (const row of data['rows']) {
        const y = row[region + ' (copies/mL)'];
        if (y !== undefined) {
            extracted.push({
                x: row['Date'],
                y: y,
                yMin: y - (row[region + ' Low Confidence Interval'] || 0),
                yMax: y + (row[region + ' High Confidence Interval'] || 0),
            });
        }
    }
    return extracted;
}

function maxExcludingOmicron(data, start, end) {
    let max = 0;
    let secondMax = 0;
    for (const row of data) {
        if (
            (start !== "" && row.x < start)
                || (end != "" && row.x > end)
                || (OMICRON_START_DATE < row.x && row.x < OMICRON_END_DATE)
        ) {
            continue;
        }
        if (row.yMax > max) {
            secondMax = max;
            max = row.yMax;
        }
    }
    // To exclude single weird outliers
    return secondMax;
}

function plot(data) {
    const northData = extractWithErrorBars(data, 'Northern');
    const southData = extractWithErrorBars(data, 'Southern');
    const northCtx = document.getElementById('northCanvas');
    const southCtx = document.getElementById('southCanvas');

    const getOptions = (region) => {
        return {
            aspectRatio: 1,
            interaction: {
                intersect: false,
            },
            scales: {
                x:  {
                    type: 'time',
                    min: startInput.value,
                    max: endInput.value,
                    time: {
                        tooltipFormat: 'MMM d, yyyy',
                    },
                },
                y: {
                    min: 0,
                    max: null,
                },
            },
            plugins: {
                title: {
                    display: true,
                    text: region + ' System Wastewater COVID RNA Signal (copies/mL)',
                },
            },
            elements: {
                pointWithErrorBar: {
                    errorBarWhiskerSize: 5,
                },
            },
        };
    };

    const northOptions = getOptions("North");
    const southOptions = getOptions("South");

    const updateYMax = () => {
        if (
            cutOmicronInput.checked
                && (startInput.value === "" || startInput.value < OMICRON_END_DATE)
                && (endInput.value === "" || endInput.value > OMICRON_START_DATE)
        ) {
            northOptions.scales.y.max = maxExcludingOmicron(northData, startInput.value, endInput.value);
            southOptions.scales.y.max = maxExcludingOmicron(southData, startInput.value, endInput.value);
        } else {
            northOptions.scales.y.max = null;
            southOptions.scales.y.max = null;
        }
    };

    updateYMax();

    const northChart = new Chart(northCtx, {
        data: {
            datasets: [
                {
                    type: 'scatterWithErrorBars',
                    label: 'Measured value',
                    data: northData,
                    backgroundColor: 'rgba(0,0,255,0.5)',
                    color: 'rgba(0,0,255,0.5)',
                    borderColor: 'rgba(0,0,255,0.5)',
                    spanGaps: true,
                },
/*                {
                    type: 'line',
                    label: 'smoothing 18/14',
                    data: smooth(northData, 'y', 18, localLinearRegression(18, 14)),
                    pointRadius: 0,
                    borderColor: 'rgba(0,0,255,0.6)',
                },
                {
                    type: 'line',
                    label: 'smoothing 14/14',
                    data: smooth(northData, 'y', 14, localLinearRegression(14, 14)),
                    pointRadius: 0,
                    borderColor: 'rgba(255,0,255,0.6)',
                },*/
                {
                    type: 'line',
                    label: '7-day average (low)',
                    data: smooth(northData, 'yMin', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(0,255,0,0.6)',
                    fill: '+1',
                    backgroundColor: 'rgba(0,255,0,0.2)',
                },
                {
                    type: 'line',
                    label: '7-day average',
                    data: smooth(northData, 'y', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(0,128,0,0.9)',
                    backgroundColor: 'rgba(0,128,0,0.9)',
                },
                {
                    type: 'line',
                    label: '7-day average (high)',
                    data: smooth(northData, 'yMax', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(0,255,0,0.6)',
                    fill: '-1',
                    backgroundColor: 'rgba(0,255,0,0.2)',
                },
            ]
        },
        options: northOptions,
    });
    const southChart = new Chart(southCtx, {
        data: {
            datasets: [
                {
                    type: 'scatterWithErrorBars',
                    label: 'Measured value',
                    data: southData,
                    backgroundColor: 'rgba(255,0,0,0.5)',
                    color: 'rgba(255,0,0,0.5)',
                    borderColor: 'rgba(255,0,0,0.5)',
                    spanGaps: true,
                },
                {
                    type: 'line',
                    label: '7-day average (low)',
                    data: smooth(southData, 'yMin', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(255,127,0,0.6)',
                    fill: '+1',
                    backgroundColor: 'rgba(255,127,0,0.2)',
                },
                {
                    type: 'line',
                    label: '7-day average',
                    data: smooth(southData, 'y', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(127,63,0,0.9)',
                    backgroundColor: 'rgba(127,63,0,0.9)',
                },
                {
                    type: 'line',
                    label: '7-day average (high)',
                    data: smooth(southData, 'yMax', 7, average),
                    pointRadius: 0,
                    borderColor: 'rgba(255,127,0,0.6)',
                    fill: '-1',
                    backgroundColor: 'rgba(255,127,0,0.2)',
                },
            ],
        },
        options: southOptions,
    });

    const update = () => {
        updateYMax();
        northChart.update();
        southChart.update();

        const params = new URLSearchParams();
        const start = startInput.value;
        if (start !== '') {
            params.set('start', start);
        }
        const end = endInput.value;
        if (end !== '') {
            params.set('end', end);
        }
        if (cutOmicronInput.checked) {
            params.set('cutOmicron', 'true');
        }
        window.location.hash = params.toString();
    };

    startInput.addEventListener('input', (e) => {
        northOptions.scales.x.min = e.target.value;
        southOptions.scales.x.min = e.target.value;
        update();
    });
    endInput.addEventListener('input', (e) => {
        northOptions.scales.x.max = e.target.value;
        southOptions.scales.x.max = e.target.value;
        update();
    });
    cutOmicronInput.addEventListener('change', (e) => {
        update();
    });
}