#!/usr/bin/env python
# -*- coding: utf-8 -*-

import yaml
import os
from DST_plot import *

from bar_plot import *


def main(system_type,date, ecvs):

    if system_type == "seasonal":
        system = "Seas5-knn"
    else:
        system = "cfsv2-knn"

    regions = ['cat']
    origdir = f'/esarchive/oper/VITIGEOSS/'
    # -------------------------------------
    # DST PLOTS 
    # -------------------------------------
    out = f'/esarchive/scratch/ptrascas/R/sunset/'
    outdir = '{}/plots/{}/{}/'.format(out,system_type,date)
    os.makedirs(outdir, exist_ok=True)

    for region in regions:

        #archive = f"/esarchive/oper/S2S-downscale/statistics/analogs/vitigeoos-{region}/"
        srcdir = f"//esarchive/oper/VITIGEOSS/output/"

        for var in ecvs:
            if system_type == "seasonal":
                leadtimes = range(3)
                archive = srcdir + f"/system51c3s/monthly_statistics/vitigeoss-{region}/"
                probs_file = f"{archive}/{var}/{date}/{var}-prob_{date}_{date[4:6]}.nc"
                skill_file = f"{archive}/{var}/{date}/{var}-skill_month{date[4:6]}.nc"
            else:
                leadtimes = range(4)
                archive = srcdir + f"/cfsv2/weekly_statistics/vitigeoss-{region}/"
                week = datetime.datetime.strptime(date, "%Y%m%d").strftime("%V")
                probs_file = f"{archive}/{var}/{date}/{var}-prob_{date}_{week}.nc"
                skill_file = f"{archive}/{var}/{date}/{var}-skill_week{week}.nc"

            for leadtime in leadtimes:

                sdate = datetime.datetime.strptime(date, "%Y%m%d")

                if system_type == "seasonal":
                    delta = datetime.timedelta(weeks=4*leadtime)
                    lead = (sdate + relativedelta(months=+leadtime+1)).strftime("%b %Y" )
                    title = 'Seasonal forecast of temperature'
                    subtitle = 'Valid for {}, issued on {}'.format(lead,sdate.strftime("%d %b %Y" ))

                else:
                    delta_days = datetime.timedelta(days=4)
                    idelta_weeks = datetime.timedelta(weeks=leadtime)
                    fdelta_weeks = datetime.timedelta(weeks=leadtime+1)
                    ileadtime = sdate + delta_days + idelta_weeks
                    fleadtime = sdate + delta_days + fdelta_weeks - datetime.timedelta(days=1)
                    title = 'Subseasonal forecast of temperature'
                    subtitle = 'Valid from {} to {}, issued on {}'.format(ileadtime.strftime("%b %d"),
                                                                          fleadtime.strftime("%b %d"),
                                                                          sdate.strftime("%d %b %Y" ))

                filename = '{}/DST-{}-{}-{}-{}-lead{}.png'.format(outdir,region,var, system, date, leadtime+1)


                with open("conf.yml") as f:
                    conf = yaml.load(f, Loader=yaml.FullLoader)
                    if var == "acprec":
                        colors = conf['conf']['colors'][var]
                    else:
                        colors = conf['conf']['colors']["default"]
                    thrs = conf['conf']['thrs']
                    locations = conf['locations'][region]

                data = get_data(probs_file,skill_file,region)

                generate_plot(data,region,locations,leadtime,thrs,colors,subtitle,filename)

    # -------------------------------------
    # BAR PLOTS 
    # -------------------------------------

    outdir = '{}/plots/{}/{}/'.format(origdir,system_type,date)

    if not os.path.exists(outdir):
        os.mkdir(outdir)
    for var in ecvs:
        for region in regions:
            if system_type == "seasonal":
                archive = f"/esarchive/oper/VITIGEOSS/output/system51c3s/monthly_statistics/vitigeoss-{region}/"
            else:
                archive = f"/esarchive/oper/VITIGEOSS/output/cfsv2/weekly_statistics/vitigeoss-{region}/"


            with open("conf.yml") as f:
                conf = yaml.load(f, Loader=yaml.FullLoader)
                locations = conf['locations'][region]

            for location,coords in locations.items():

                data = get_bardata(var,date,archive,location,coords,region,system_type)
                terciles = data[['p33','p66']]
                extremes = data[['p10','p90']]
                skills = data[['rpss','bsp90','bsp10']]
                data = df_to_pd(data)

                if var == "acprec":
                    colors = load_colors("conf.yml",var)
                else:
                    colors = load_colors("conf.yml","default")

                var_attrs = load_var_attrs("conf.yml",var)
                if system_type == "seasonal":
                    bar_plot(var,date,location,data,colors,var_attrs,skills,terciles,extremes,outdir)
                else:
                    bar_plot_subseas(var,date,location,data,colors,var_attrs,skills,terciles,extremes,outdir)


if __name__ == "__main__":

    sys_type = sys.argv[1]
    #run_conf = "conf/vitigeoos-seasonal-run.yml"
    run_sdate = sys.argv[2]
    target_var = sys.argv[3]
    #run_sdate = "20210701"

    main(sys_type, run_sdate, [target_var])


