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,actualfitness=0, actualfatigue=0):
    fatigue = actualfatigue
    fitness = actualfitness
    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, reset=False):

    try:
        w_last = Workout.objects.get(id=rower.last_workout)
        datetime_last = w_last.startdatetime
    except Workout.DoesNotExist:
        w_last = None
        datetime_last = datetime.datetime.now() - datetime.timedelta(days=90)

    if reset:
        rower.last_workout = 0
        rower.actualfit = 0
        rower.actualfatigue = 0
        rower.save()
        datetime_last = datetime.datetime.now() - datetime.timedelta(days=90)

    if not datetime_last:
        datetime_last = datetime.datetime.now() - datetime.timedelta(days=90)

    workouts = Workout.objects.filter(user=rower, startdatetime__gt=datetime_last,
                                      duplicate=False).order_by('date').exclude(id=rower.last_workout)

    kfitness = rower.kfit
    kfatigue = rower.kfatigue
    actualfitness = rower.actualfit
    actualfatigue = rower.actualfatigue

    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 actualfitness, actualfatigue

    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,
        actualfitness=actualfitness,
        actualfatigue=actualfatigue,
    )

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

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

    if workouts:
        enddate = workouts.last().date
        rower.last_workout = workouts.last().id
        rower.actualfit = int(endfitness)
        rower.actualfatigue = int(endfatigue)
        rower.save()

    return endfitness, endfatigue