#! /usr/bin/env python
# Copyright 2022 John Hanley. MIT licensed.
from pathlib import Path

from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import typer

from problem.covid.us_cases_deaths import get_cases_and_deaths


def _get_daily_cases_and_deaths():
    df = get_cases_and_deaths()
    df.cases = df.cases.diff().shift(-1)
    df.deaths = df.deaths.diff().shift(-1)
    df = df.dropna()
    df.cases = df.cases.astype(np.int32)
    df.deaths = df.deaths.astype(np.int32)

    # clip a March 2022 negative deaths figure
    df.cases = df.cases.clip(lower=0)
    df.deaths = df.deaths.clip(lower=0)

    df = df.set_index('date')
    return df


def predict(out_file=Path('~/Desktop/lag.png')):
    out_file = Path(out_file).expanduser()

    df = _get_daily_cases_and_deaths()
    train, test = _split(df)
    assert len(train) >= 224

    # dates = df.index.get_level_values(0)
    # from sklearn.preprocessing import StandardScaler
    # x_train = np.array(StandardScaler().fit_transform(train.cases))
    x_train = np.array(train.cases)
    y_train = np.array(train.deaths)
    model = LinearRegression()
    model.fit(x_train.reshape(-1, 1), y_train)

    x_test = np.array(test.cases)
    y_test = np.array(test.deaths)
    y_pred = model.predict(x_test.reshape(-1, 1))

    _, axs = plt.subplots(1, 2)
    plt.sca(axs[0])
    train['cases'] /= 1e2
    sns.scatterplot(data=train)
    plt.xticks(rotation=45)
    y_limit = 6_000
    plt.gca().set_ylim((0, y_limit))

    plt.sca(axs[1])
    test['cases'] /= 1e2
    test['pred_deaths'] = y_pred
    sns.scatterplot(data=test)
    plt.xticks(rotation=45)
    plt.gca().set_ylim((0, y_limit))
    plt.savefig(out_file)
    plt.show()

    print(f'R2: {model.score(x_test.reshape(-1, 1), y_test)}')
    print(f'MSE: {np.mean((y_test - y_pred)**2)}')
    print(f'MAE: {np.mean(np.abs(y_test - y_pred))}')


def _split(df: pd.DataFrame, split_date='2020-09-01'):
    train = df.query(f'date < "{split_date}"')
    test = df.query(f'date >= "{split_date}"')
    assert len(df) == len(train) + len(test)  # no rows left behind
    return (train.copy(),
            test.copy())


if __name__ == '__main__':
    typer.run(predict)

@@ -103,44 +103,190 @@
 arn.
-linear_model import LinearRegression
+ensemble import RandomForestRegressor%0Afrom sklearn.linear_model import LinearRegression, SGDRegressor%0Afrom sklearn.svm import LinearSVR%0Afrom sklearn.tree import DecisionTreeRegressor
 %0Aimp
@@ -1311,24 +1311,25 @@
 ain.deaths)%0A
+%0A
     model = 
@@ -1328,32 +1328,37 @@
 model = 
-Linear
+RandomForest
 Regress
-ion
+or
 ()%0A    m
@@ -1919,23 +1919,8 @@
 ile)
-%0A    plt.show()
 %0A%0A  
@@ -1976,16 +1976,20 @@
  y_test)
+:.3f
 %7D')%0A    
@@ -2031,16 +2031,20 @@
 red)**2)
+:.3f
 %7D')%0A    
@@ -2089,16 +2089,20 @@
 y_pred))
+:.3f
 %7D')%0A%0A%0Ade

Use RandomForestRegressor