extern crate chrono;
extern crate csv;
extern crate ini;

use chrono::prelude::*;
use chrono::Utc;

mod curves {
    mod exponential;
}
pub mod estimators {
    pub mod ratios;
    pub mod exponential;
    pub mod gompertz;
}

mod datamodel;
mod data;

pub mod search {
    pub mod ternary;
    pub mod golden_section_search;
}

use estimators::ratios::{estimate_growth_rate};
use estimators::exponential::estimate_exponential;
use estimators::gompertz::{estimate_gompertz_with_recovery};
use search::golden_section_search::golden_section_search;

fn main() -> std::io::Result<()> {
    let saturation_percentage = 66.0;
    let population = 83019213;
    let plateau = population as f64 / 100.0 * saturation_percentage;
    let days = crate::data::berliner_morgenpost_deutschland();
    /*
     * https://www.imperial.ac.uk/media/imperial-college/medicine/sph/ide/gida-fellowships/Imperial-College-2019-nCoV-severity-10-02-2020.pdf
     */
    let recovery_period = 12.0;

    let gompertz_reconstruction = |rate: f64| {
        let initial_infected = days[0].infected as f64;
        let mut ret = Vec::new();
        let mut day: f64 = 0.0;

        for _ in &days {
            ret.push(estimate_gompertz_with_recovery(initial_infected, rate, plateau, recovery_period, day));
            day = day + 1.0;
        }

        ret
    };

    let gompertz_squared_errors = |rate: f64| {
        let mut ret = Vec::new();

        let reconstruction = gompertz_reconstruction(rate);
        let mut actual = Vec::new();
        for each in &days {
            actual.push(each.infected as f64);
        }
        for (a, b) in reconstruction.iter().zip(actual.iter()) {
            ret.push((a - b).powi(2));
        }

        ret
    };

    let gompertz_root_mean_squared_error = |rate: f64| -> f64 {
        let errors = gompertz_squared_errors(rate);
        let sample_size = errors.len() as f64;
        let sum: f64 = errors.iter().sum();

        (sum / sample_size).sqrt()
    };

    let find_gompertz_rate = |start: f64, end:f64| -> f64 {
        golden_section_search(&gompertz_root_mean_squared_error, start, end, 0.00000000000001)
    };

    let gompertz_rate = find_gompertz_rate(0.0, 0.1);
    let gompertz_squared_error = gompertz_root_mean_squared_error(gompertz_rate);
    println!("gompertz rate: {}, rms: {}", gompertz_rate, gompertz_squared_error.sqrt());

    let mean = estimate_growth_rate(&days);

    println!("estimated growth rate: {:.0}%±{:.0}%\n\n", (mean.geometric_mean * 100.0).round(), (mean.geometric_mean_stderr * 100.0).round());

    use std::fs::File;

    let buffer = File::create("corona.csv")?;
    let mut writer = csv::Writer::from_writer(buffer);

    for day in 0 .. 400 {
        let day = day as f64;
        let initial_infected = days[0].infected as f64;
        let estimate_exponential = |rate| { estimate_exponential(initial_infected, rate, day) };
        let rate = mean.geometric_mean;

        let exponential_estimate             = estimate_exponential(rate);

        let gompertz_estimate = estimate_gompertz_with_recovery(initial_infected, gompertz_rate, plateau, recovery_period, day).round();

        let mut write_record = |date: Date<Utc>| {
            let date = date.format("%Y-%m-%d").to_string();
            let exponential_estimate = exponential_estimate.to_string();
            let gompertz_estimate = gompertz_estimate.to_string();
            let actual_infected = if (day as usize) < days.len() {
                days[day as usize].infected.to_string()
            } else {
                "".to_string()
            };

            writer.write_record(&[date, actual_infected, exponential_estimate, gompertz_estimate]).unwrap();
        };

        if (day as usize) < days.len() {
            let date = days[day as usize].day;
            write_record(date);
        } else {
            let date = days[days.len() - 1].day;
            let date = date.checked_add_signed(chrono::Duration::days(day as i64 - days.len() as i64 + 1)).unwrap();

            write_record(date);
        }
    }

    let mut conf = ini::Ini::new();
    let last_day_index = days.len() - 1;
    let infected = days[last_day_index].infected;
    let dead = days[last_day_index].dead;
    let mortality_rate_percent = 0.6;
    let severe_cases_rate_percent = 5.0;
    // https://www.medicalnewstoday.com/articles/coronavirus-81-of-cases-are-mild-study-says
    let medically_relevant_cases_percent = 19.1;
    let available_icus = 28000 + 8000;
    let available_beds = 497000;
    conf.with_section(Some("most_recent_values"))
        .set("infected", infected.to_string())
        .set("dead", dead.to_string())
        .set("current_day", last_day_index.to_string())
        .set("data_source", "Berliner Morgenpost");
    conf.with_section(Some("assumptions"))
        .set("mortality_rate_percent", mortality_rate_percent.to_string())
        .set("severe_cases_rate_percent", severe_cases_rate_percent.to_string())
        .set("medically_relevant_cases_percent", medically_relevant_cases_percent.to_string())
        .set("available_icus", available_icus.to_string())
        .set("available_beds", available_beds.to_string())
        .set("recovery_period", recovery_period.to_string());
    conf.with_section(Some("exponential"))
        .set("rate", mean.geometric_mean.to_string())
        .set("rate_stderr", mean.geometric_mean_stderr.to_string());
    conf.with_section(Some("gompertz"))
        .set("rate", gompertz_rate.to_string())
        .set("rmserror", gompertz_squared_error.sqrt().to_string())
        .set("saturation_percentage", saturation_percentage.to_string())
        .set("population", population.to_string());
    conf.write_to_file("estimates.ini").unwrap();

    Ok(())
}