// SARS-CoV-2-Viz
// Animated COVID case count visualization
// Copyright 2022 Edward L. Platt <ed@elplatt.com>

const covidUrl = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_confirmed_US.csv";

const usUrl = "https://cdn.jsdelivr.net/npm/us-atlas@3/counties-albers-10m.json";

const metadataUrl = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/UID_ISO_FIPS_LookUp_Table.csv";

async function fetchWithProgress(url, onProgress) {

    // Make request and get headers
    let response = await fetch(url);
    const reader = response.body.getReader();
    const contentLength = response.headers.get('Content-Length');

    // Receive chunks
    let chunks = [];
    let contentReceived = 0;
    while (true) {
	const {done, value} = await reader.read()
	if (done) { break; }
	chunks.push(value);
	let bytes = value.length;
	contentReceived += bytes;

	// Update progress
	onProgress(contentReceived);
    }

    // Combine chunks and decode
    let contentBytes = new Uint8Array(contentReceived);
    let position = 0;
    chunks.forEach((chunk) => {
	contentBytes.set(chunk, position);
	position += chunk.length;
    });
    let content = new TextDecoder("utf-8").decode(contentBytes);
    return content;
}

// Convert csv to tsv
// Assumes quoted strings do not contain escaped quotes
function csvToTSV(csv, onProgress) {
    let quoteParts = csv.split('"');
    let tsv = "";
    for (const [index, part] of quoteParts.entries()) {
	if (index % 2 == 0) {
	    // Not quoted
	    tsv += part.replace(/,/g, "\t");
	} else {
	    // Quoted
	    tsv += part;
	}
    }
    return tsv;
}

function tsvRowToJSON(keys, tsvRow) {
    let row = {};
    let rowParts = tsvRow.split("\t");
    for (const [index, key] of keys.entries()) {
	row[key] = rowParts[index];
    }
    return row;
}

// Convert tsv to json
function parseCovidData(tsv, onProgress) {
    let rows = tsv.split(/[\r\n]+/);
    let keys = rows[0].split("\t");
    let dates = keys.slice(11);

    // Initialize data array
    let dataForDay = [];
    for (const date of dates) {
	dataForDay.push([]);
    }
    
    // Iterate through non-header rows
    const rowCount = rows.length - 1;
    let percent = 0;
    for (const [index, row] of rows.slice(1).entries()) {
	const rowData = tsvRowToJSON(keys, row);

	// Only update progress when it has actually changed
	let newPercent = Math.floor(100 * index / rowCount);
	if (newPercent > percent) {
	    percent = newPercent;
	    onProgress(`${percent}%`);
	}   
	
	let FIPS = rowData.FIPS;
	if (typeof(FIPS) === "undefined") {
	    continue;
	}
	FIPS = FIPS.slice(0, FIPS.indexOf(".") || FIPS.length).padStart(5, "0");

 	let lastSeven = [];
	for (const [index, date] of dates.entries()) {
	    
	    const newCases = index == 0 ?
		  parseInt(rowData[date]) :
		  parseInt(rowData[date])
		  - parseInt(rowData[dates[index - 1]]);

	    // Update array for running average
	    if (lastSeven.length == 7) {
		lastSeven.shift();
	    }
	    lastSeven.push(newCases);
	    const mean = Math.round(
		lastSeven.reduce((x, y) => x + y) / lastSeven.length);
	    
	    dataForDay[index].push({
		FIPS: FIPS,
		date: date,
		count: newCases,
		sevenDayMean: mean
	    });
	}
    }
    onProgress("100%");
    
    // Calculate maximum case count for each day
    let maxSoFar = 0;
    for (const [day, date] of dates.entries()) {

	let dayMax = Math.max(
	    ...dataForDay[day].map(
		d => d.sevenDayMean));

	maxSoFar = Math.max(maxSoFar, dayMax);

	// Add data to each county for current day
	for (let [countyIndex, d] of dataForDay[day].entries()) {
	    d.dayMax = dayMax;
	    d.maxSoFar = maxSoFar;
	}
    }

    return dataForDay.map(data => {
	const entries = new Map(data.map(d => [d.FIPS, d]));
	return Object.fromEntries(entries);
    });
}

async function getData(onProgress) {
    const content = await fetchWithProgress(covidUrl, onProgress);
    return content;
}

async function parseData(content, onProgress) {
    let tsv = csvToTSV(content);
    let dataForDay = parseCovidData(tsv, onProgress);
    return dataForDay;
}

async function getFeatures(onProgress) {
    const content = await fetchWithProgress(usUrl, onProgress);
    return content;
}

async function parseFeatures(content, onProgress) {
    onProgress("0%");
    let features = JSON.parse(content);
    onProgress("100%");
    return features;
}

async function getMetadata(onProgress) {
    const content = await fetchWithProgress(metadataUrl, onProgress);
    return content;
}

function tsvToJSON(tsv, onProgress) {
    const lines = tsv.split(/\n/);
    let columnNames = [];
    let data = [];
    const count = lines.length;
    for (const [row, line] of lines.entries()) {
	if (row == 0) {
	    columnNames = line.split("\t").map((x) => x.trim());
	} else {
	    let dataRow = {};
	    for (const [col, cell] of line.split("\t").entries()) {
		dataRow[columnNames[col]] = cell.trim();
	    }
	    // Limit to US Counties
	    if (dataRow["iso3"] == "USA"
		&& dataRow["Admin2"] != ""
		&& dataRow["Admin2"] != "Unassigned"
		&& dataRow["FIPS"] != "") {
		data.push(dataRow);
	    }
	}
	let progress = count / row;
	onProgress(progress);
    }
    return data;
}

function parseMetadataTSV(tsv, onProgress) {
    return tsvToJSON(tsv, onProgress);
}

async function parseMetadata(content, onProgress) {
    onProgress("0%");
    let tsv = csvToTSV(content);
    let metadata = parseMetadataTSV(tsv, onProgress);
    onProgress("100%");
    return metadata;
}

function createChoropleth(us, msPerFrame) {

    const counties = topojson.feature(us, us.objects.counties);
    const states = topojson.feature(us, us.objects.states)
    const statemesh = topojson.mesh(
	us, us.objects.states, (a, b) => a !== b)
    const statemap = new Map(states.features.map(d => [d.id, d]));

    let chart = Choropleth({
	id: d => d.FIPS,
	value: d => d.sevenDayMean,
	domain: d => {
	    let domain = [0, Math.max(1, d["00000"].dayMax)];
	    return domain;
	},
	title: (f, d) => `${f.properties.name}, ${statemap.get(f.id.slice(0, 2)).properties.name}\n${d?.sevenDayMean}`,
	features: counties,
	borders: statemesh,
	width: 975,
	height: 610,
	duration: msPerFrame
    });

    return chart;
}

function UI() {

    let frames = [];
    let currentFrame = 0;
    let msPerFrame = 500;
    let playing = false;
    let onTick = null;

    let animate = function () {
	if (playing) {
	    setTimeout(animate, msPerFrame);
	}
	if (onTick) {
	    onTick(currentFrame);
	}
	currentFrame = (currentFrame + 7) % frames.length;
    };
    
    let node = document.createElement("div");
    
    let controls = document.createElement("div");
    node.appendChild(controls);
   
    let play = document.createElement("button");
    play.innerText = 'play';
    play.id = 'play';
    controls.appendChild(play);
    play.addEventListener(
	"click", () => {
	    if (playing) {
		play.innerText = 'play';
		playing = false;
	    } else {
		play.innerText = 'pause';
		playing = true;
		animate();
	    }
	});

    let display = document.createElement("div");
    display.id = "display";
    let mainDisplay = document.createElement("div");
    mainDisplay.id = "main-display";
    let secondaryDisplays = {};
    display.appendChild(mainDisplay);
    node.appendChild(display);

    return {
	node: node,
	display: (s, id=null) => {
	    if (id === null) {
		mainDisplay.innerText = s;
	    } else {
		secondaryDisplays[id].innerText = s;
	    }
	},
	addDisplay: (id) => {
	    let secondary = document.createElement("div");
	    secondary.id = `display-${id}`;
	    display.appendChild(secondary);
	    secondaryDisplays[id] = secondary;
	},
	clearDisplays: () => {
	    mainDisplay.innerText = '';
	    for (const [id, d] of Object.entries(secondaryDisplays)) {
		d.remove();
	    }
	    secondaryDisplays = {};
	},
	onTick: (f) => { onTick = f; },
	setFrames: (newFrames, current) => {
	    frames = newFrames;
	    currentFrame = current;
	},
	setMsPerFrame: (newMsPerFrame) => { msPerFrame = newMsPerFrame; }
    }
}

function ProgressTracker(onCreate, onProgress) {
    let jobProgress = [];
    let tracker = {
	getTracker: (initial='') => {
	    let id = jobProgress.length;
	    jobProgress.push(initial);
	    onCreate(id);
	    return (progress) => {
		jobProgress[id] = progress;
		onProgress(id, tracker.progress(id));
	    }
	},
	progress: (id) => jobProgress[id]
    };
    return tracker;
}

function main() {
    let ui = UI();
    let tracker = ProgressTracker(
	(id) => { ui.addDisplay(id); },
	(id, progress) => ui.display(progress, id)
    );
    document.getElementById("header").appendChild(ui.node);
    
    let onDataProgress = tracker.getTracker();
    let dataPromise = getData(
	(progress) => {
	    onDataProgress(`Fetching data: Received ${progress} bytes`)
	})
	.then((dataCSV) => {
	    return parseData(
		dataCSV,
		(progress) => onDataProgress(`Parsing data: ${progress}`)
	    );
	});

    let onFeatureProgress = tracker.getTracker();
    let featurePromise = getFeatures(
	(progress) => {
	    onFeatureProgress(`Fetching features: Received ${progress} bytes`);
	})
	.then((features) => {
	    return parseFeatures(
		features,
		(progress) => onFeatureProgress(`Parsing features: ${progress}`));
	});

    let onMetadataProgress = tracker.getTracker();
    let metadataPromise = getMetadata(
	(progress) => {
	    onMetadataProgress(`Fetching metadata: ${progress}`);
	})
	.then((metadata) => {
	    return parseMetadata(
		metadata,
		(progress) => onMetadataProgress(`Parsing metadata: ${progress}`));
	});
    
    Promise.all([
	dataPromise, featurePromise, metadataPromise
    ]).then((values) => {
	const [dataForDay, us, metadata] = values;
	console.log(metadata);
	ui.clearDisplays();
	ui.display("Initializing visualization");
	let msPerFrame = 200;
	console.log(dataForDay);
	console.log(us);
	console.log(metadata);
	let dates = dataForDay.map((data) => data["00000"].date);
	let chart = createChoropleth(us, msPerFrame);
	document.getElementById("content").appendChild(chart.node);
	let day = dataForDay.length - 1;
	let data = dataForDay[day];
	ui.display(dataForDay[day]["00000"].date);
	ui.setFrames(dates, day);
	ui.setMsPerFrame(msPerFrame);
	chart.update(data);
	ui.onTick((day) => {
	    chart.update(dataForDay[day]);
	    ui.display(dates[day]);
	});
    });
};

main();