import polars as pl
import datetime
import arrow
import numpy as np

from rowers.dataprep import rscore_approx

from rowers.models import Rower, Workout

def weightfromrecord(row,metricchoice):
    vv = row[metricchoice][0]
    if vv > 0:
        return vv
    if metricchoice == 'rscore': # pragma: no cover
        return rscore_approx(row)

    return 0 # pragma: no cover


def getfatigues(df, fatigues, fitnesses, dates, impulses, startdate, kfatigue, kfitness):
    fatigue = 0
    fitness = 0
    impulses = []

    for f in fatigues:
        impulses.append(0)

    lambda_a = 2/(kfatigue+1)
    lambda_c = 2/(kfitness+1)

    enddate = datetime.date.today()
    nrdays = (enddate - startdate).days

    for i in range(nrdays+1):
        date = startdate + datetime.timedelta(days=i)
        datekey = date.strftime('%Y-%m-%d')
        weight = 0

        try:
            df2 = df.filter(pl.col('date') == date)
            if type(df2) == pl.Series:
                weight += weightfromrecord(df2, 'rscore')
            else:
                for row in df2.iter_slices(n_rows=1):
                    weight += weightfromrecord(row, 'rscore')
        except KeyError:
            pass

        impulses.append(weight)

        fatigue = (1 - lambda_a) * fatigue + lambda_a * weight
        fitness = (1 - lambda_c) * fitness + lambda_c * weight

        fatigues.append(fatigue)
        fitnesses.append(fitness)
        dates.append(arrow.get(date).datetime)

    return fatigues, fitnesses, dates, impulses

def calculate_fitness(rower):
    workouts = Workout.objects.filter(user=rower, date__gte=datetime.date.today() - datetime.timedelta(days=90),
                                      duplicate=False).order_by('date')

    kfitness = rower.kfit
    kfatigue = rower.kfatigue
    

    records = []

    for w in workouts:
        dd = {
            'date': w.date,
            'trimp': w.trimp,
            'rscore': w.rscore,
            'hrtss':w.hrtss,
            'duration': w.duration,
            'id': w.id,
            'rpe': w.rpe,
        }
        records.append(dd)

    df = pl.from_records(records)
    if df.is_empty():
        return 0, 0

    dates = [arrow.get(w.date).datetime for w in workouts]
    impulses = [np.nan for w in workouts]
    fitnesses = [np.nan for w in workouts]
    fatigues = [np.nan for w in workouts]

    fatigues, fitnesses, dates, impulses = getfatigues(
        df,
        fatigues,
        fitnesses,
        dates,
        impulses,
        datetime.date.today()-datetime.timedelta(days=90),
        kfatigue, kfitness
    )

    df = pl.DataFrame({
        'date': dates,
        'fatigue': fatigues,
        'fitness': fitnesses,
        'impulse': impulses,
    })

    endfitness = fitnesses[-2]
    endfatigue = fatigues[-2]

    return endfitness, endfatigue