#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Cloud HR experiments

Created on Fri Dec  9 11:22:51 2016

@author: maxwell
"""

import numpy as np
import matplotlib.pyplot as plt
import time
import copy

import atmosphere as a
from parm import ChemParm, LWParm, SWParm
from solver import SolverFactory,HR
from misc.humidity import manaberh
import misc.solargeometry as solar

# %% set up
st = time.clock()
timestr = "Elapsed Time: {:4f}s"

plev =np.logspace(-2,np.log10(1013), 601)
lat = 10
decl = 21.11
rhlabel= ''
holdrh=True
rh = np.ones(len(plev))*0.5
holdtsfc=True
cpdair=1004.0

timestep=0.25
radmodel='fu'
gridstagger=False
maxsteps=300
tol = .011

cldprofiles=[ None,
             'atmosphere/profiles/hr/yang/cf_allclouds/annual_mean_20Nto20S_cf.txt',
             'atmosphere/profiles/hr/yang/cf_withoutcirrus/annual_mean_20Nto20S_cf_withoutcirrus.txt',
              None,
             'atmosphere/profiles/hr/yang/cf_allclouds/annual_mean_wp_cf.txt',
             'atmosphere/profiles/hr/yang/cf_withoutcirrus/annual_mean_wp_cf_withoutcirrus.txt'
              ]
o3profiles=['atmosphere/profiles/ozone/yang/annual_ozone_20Nto20S.dat',
            'atmosphere/profiles/ozone/yang/annual_ozone_20Nto20S.dat',
            'atmosphere/profiles/ozone/yang/annual_ozone_20Nto20S.dat',
            'atmosphere/profiles/ozone/yang/annual_ozone_fiji.dat',
            'atmosphere/profiles/ozone/yang/annual_ozone_fiji.dat',
            'atmosphere/profiles/ozone/yang/annual_ozone_fiji.dat'
             ]

anames=['Trop', 'Trop CF', 'Trop CF (no Ci)', 'WP', 'WP CF', 'WP CF (no Ci)']
cldconds = dict(zip(anames, cldprofiles))
o3conds = dict(zip(anames, o3profiles))
ts = 300.0
nsim = len(anames)
atms=dict.fromkeys(anames)
atms_noco2 = dict.fromkeys(anames)
atms_noo3 = dict.fromkeys(anames)
atms_wvonly = dict.fromkeys(anames)
hr=dict.fromkeys(anames)
auxhr=dict.fromkeys(anames)
hr_noco2=dict.fromkeys(anames)
hr_noo3 = dict.fromkeys(anames)
hr_wvonly= dict.fromkeys(anames)
flx = dict.fromkeys(anames)
flx_noco2 = dict.fromkeys(anames)
flx_noo3 = dict.fromkeys(anames)
flx_wvonly = dict.fromkeys(anames)

swparm = dict.fromkeys(anames)

cparm = ChemParm()
cparm_noco2=ChemParm(co2ppmv=1.0e-4)
lwparm = LWParm()
slv = SolverFactory.create(kind='rce', timestep=timestep, holdtsfc=holdtsfc,
                           radmodel=radmodel,cpdair=cpdair,tol=tol,
                           maxsteps=maxsteps)
radslv = SolverFactory.create(kind='rad',radmodel=radmodel,
                              cpdair=cpdair)

st2 = st
for name,cond in cldconds.items():
    print('SOLVING {}'.format(name))
    prof = 'jtrp'
    atms[name] = a.Atmosphere.mcclatchy(prof, p=plev, rhlev=rh, holdrh=holdrh,
                                        gridstagger=gridstagger,tsfc=ts)
    atms[name].ozone_fromfile(o3conds[name])
    atms_noco2[name] = copy.deepcopy(atms[name])
    atms_noo3[name] = copy.deepcopy(atms[name])
    atms_noo3[name].o3 = np.zeros(len(atms_noo3[name]))
    atms_wvonly[name] = copy.deepcopy(atms[name])
    atms_wvonly[name].o3 = np.zeros(len(atms_wvonly[name]))

    mu=solar.mubar(lat,decl)
    fday=solar.fday(lat,decl)
    swparm[name] = SWParm(coszen=mu,fday=fday)

    if cond is not None:
        auxhr[name]=HR.fromfile(atms[name], cond)
    else:
        zeros = np.zeros(len(atms[name]))
        auxhr[name]=HR(zeros, zeros)
    slv.auxhr=auxhr[name]

    atms[name],flx[name],hr[name] = slv.solve(
                                        atms[name],cparm,lwparm,swparm[name])
    atms_noco2[name], flx_noco2[name], hr_noco2[name] = radslv.solve(
        atms_noco2[name],cparm_noco2, lwparm,swparm[name])
    atms_noo3[name], flx_noo3[name], hr_noo3[name] = radslv.solve(
        atms_noo3[name],cparm,lwparm,swparm[name])
    atms_wvonly[name], flx_wvonly[name], hr_wvonly[name] = radslv.solve(
        atms_wvonly[name], cparm_noco2, lwparm,swparm[name])
    ed = time.clock()
    print(timestr.format(ed-st2))
    st2 = ed

print('Total time: {}'.format(ed-st))

# %% plot temps

plt.figure(1)
plt.clf()
yls = (8,1013)
#ytcks = [10,20,40,60,80,100,200,400,600,800,1000]
xls = (150,300)
xtcks = np.linspace(150,300,3)
plt.suptitle('Temperature (K)')

if(gridstagger):
    fstag='0'
else:
    fstag='1'

for i, name in enumerate(anames):
    ax = plt.subplot(1,nsim,i+1)
    plt.hold('on')
    try:
        plt.semilogy(atms[name].t, atms[name].p)
        plt.plot(atms[name].tconv, atms[name].pconv, 'ks')
        plt.plot(atms[name].tcold, atms[name].pcold, 'ko')
    except ValueError:
        pass
    finally:
        plt.ylim(yls)
    #    plt.yticks(ytcks)
        plt.xlim(xls)
        plt.xticks(xtcks)
        plt.xlabel(name.upper())
        ax.invert_yaxis()
        if (i==0):
            plt.ylabel('Pressure (hPa)')
    #        ax.yaxis.set_ticklabels(['{:.0f}'.format(tick) for tick in ytcks])
        else:
            ax.yaxis.set_ticklabels([])
else:
    plt.show()

figname = 'img/rce_{}g{}_cld_{}_eq.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% plot hr (net)
plt.figure(2)
plt.clf()
yls = (10,1013)
xls = (-3,2)
xtcks = np.linspace(xls[0],xls[1],2)
#ytcks = np.linspace(,1000,10)
plt.suptitle('HR (K/d)')

for i, name in enumerate(anames):
    ax = plt.subplot(1,nsim,i+1)
    plt.hold('on')
    iconv = atms[name].iconv
    icold = atms[name].icold
    try:
        plt.semilogy(hr[name].hrir, atms[name].p,color='crimson')
        plt.semilogy(hr[name].hrsw, atms[name].p,'mediumblue')
        plt.semilogy(hr[name].hr, atms[name].p,'orange')
        plt.plot(np.zeros(len(atms[name])), atms[name].p, 'k--')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pcold, 'k')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pconv, 'k')

    except ValueError:
        pass
    finally:
        plt.ylim(yls)
    #    plt.yticks(ytcks)
        plt.xlim(xls)
        plt.xticks(xtcks)
        plt.xlabel(name.upper())
        ax.invert_yaxis()

        if (i==0):
            plt.ylabel('Pressure (hPa)')
    #        ax.yaxis.set_ticklabels(['{:.0f}'.format(tick) for tick in ytcks])
        else:
            ax.yaxis.set_ticklabels([])
else:
    plt.show()

figname = 'img/rce_{}g{}_cld_{}_hr.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% co2-only hr
plt.figure(3)
plt.clf()
yls = (10,1013)
xls = (-3,3)
xtcks = np.linspace(xls[0],xls[1],3)
#ytcks = np.linspace(100,1000,10)
plt.suptitle('CO2 IR HR (K/d)')

for i, name in enumerate(anames):
    ax = plt.subplot(1,nsim,i+1)
    plt.hold('on')
    iconv = atms[name].iconv
    icold = atms[name].icold

    try:
        plt.plot(hr[name].hrir - hr_noco2[name].hrir, atms[name].p)
        plt.plot(np.zeros(len(atms[name])), atms[name].p, 'k--')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pcold, 'k')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pconv, 'k')

    except ValueError:
        pass
    finally:
        plt.ylim(yls)
        plt.xlim(xls)
        plt.xticks(xtcks)
        plt.xlabel(name.upper())
        ax.invert_yaxis()
        ax.set_yscale('log')
    #    plt.yticks(ytcks)

        if (i==0):
            plt.ylabel('Pressure (hPa)')
    #        ax.yaxis.set_ticklabels(['{:.0f}'.format(tick) for tick in ytcks])
        else:
            ax.yaxis.set_ticklabels([])
else:
    plt.show()

figname = 'img/rce_{}g{}_cld_{}_hrco2.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% o3-only hr
plt.figure(4)
plt.clf()
yls = (10,1013)
xls = (-1,2)
xtcks = np.linspace(xls[0],xls[1],4)
#ytcks = np.linspace(100,1000,10)
plt.suptitle(' O3 HR (K/d)')

for i, name in enumerate(anames):
    ax = plt.subplot(1,nsim,i+1)
    plt.hold('on')
    iconv = atms[name].iconv
    icold = atms[name].icold
    try:
#        plt.plot(hr[name].hrir - hr_noo3[name].hrir, atms[name].p,color='crimson')
#        plt.plot(hr[name].hrsw - hr_noo3[name].hrsw, atms[name].p,color='mediumblue')
        plt.plot(hr[name].hr - hr_noo3[name].hr, atms[name].p,'orange')
        plt.plot(np.zeros(len(atms[name])), atms[name].p, 'k--')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pcold, 'k')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pconv, 'k')
    except ValueError:
        pass
    finally:
        plt.ylim(yls)
        plt.xlim(xls)
        plt.xticks(xtcks)
        plt.xlabel(name.upper())
        ax.invert_yaxis()
        ax.set_yscale('log')
    #    plt.yticks(ytcks)

        if (i==0):
            plt.ylabel('Pressure (hPa)')
    #        ax.yaxis.set_ticklabels(['{:.0f}'.format(tick) for tick in ytcks])
        else:
            ax.yaxis.set_ticklabels([])
else:
    plt.show()

figname = 'img/rce_{}g{}_cld_{}_hro3.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% wv-only hr
plt.figure(5)
plt.clf()
yls = (10, 1013)
xls = (-2, 2)
xtcks = np.linspace(-2,2,3)
plt.suptitle('H2O HR')

for i, name in enumerate(anames):
    ax = plt.subplot(1,nsim,i+1)
    plt.hold('on')
    iconv = atms[name].iconv
    icold = atms[name].icold
    try:
#        plt.plot(hr[name].hrir - hr_noo3[name].hrir, atms[name].p,color='crimson')
#        plt.plot(hr[name].hrsw - hr_noo3[name].hrsw, atms[name].p,color='mediumblue')
        plt.plot(hr_wvonly[name].hr, atms[name].p,'g')
        plt.plot(np.zeros(len(atms[name])), atms[name].p, 'k--')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pcold, 'k')
        plt.plot(xtcks, np.ones(len(xtcks))*atms[name].pconv, 'k')
    except ValueError:
        pass
    finally:
        plt.ylim(yls)
        plt.xlim(xls)
        plt.xticks(xtcks)
        plt.xlabel(name.upper())
        ax.invert_yaxis()
        ax.set_yscale('log')
    #    plt.yticks(ytcks)

        if (i==0):
            plt.ylabel('Pressure (hPa)')
    #        ax.yaxis.set_ticklabels(['{:.0f}'.format(tick) for tick in ytcks])
        else:
            ax.yaxis.set_ticklabels([])
else:
    plt.show()

figname = 'img/rce_{}g{}_cld_{}_hrwv.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% compare o3 profiles
plt.figure(6)
plt.clf()
xls1 = (0,18)
xls2 = (-1.6, 0.2)
yls = (1, 1013)

plt.subplot(1,2,1)
plt.hold('on')
plt.plot(atms['Trop'].o3*1e6,atms['Trop'].p, label='All Tropics')
plt.plot(atms['WP'].o3*1e6, atms['WP'].p, label='WP')
plt.ylim(yls)
plt.xlim(xls1)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('O3 Profile')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Conc. ($10^{-6}$ g/g)')
plt.legend(loc='best')

plt.subplot(1,2,2)
plt.hold('on')
plt.plot(1e6*(atms['WP'].o3-atms['Trop'].o3), atms['WP'].p)
plt.plot(np.zeros(len(atms['WP'])), atms['WP'].p, 'k--')
plt.ylim(yls)
plt.xlim(xls2)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('$\Delta$ O3')
plt.xlabel('Conc. ($10^{-6}$ g/g)')

figname = 'img/rce_{}g{}_cld_{}_o3prof.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% show delta T for clouds
plt.figure(7)
plt.clf()
xls1 = (-10,10)
xls2 = (-10, 10)
yls = (40, 400)

plt.subplot(1,2,1)
plt.hold('on')
names = anames[1:3]
ref = anames[0]
for name in names:
    plt.plot(atms[name].t-atms[ref].t,atms[name].p,label=name)
    plt.plot(atms[name].tconv-atms[ref].t[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(atms[name].tcold-atms[ref].t[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
plt.ylim(yls)
plt.xlim(xls1)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Temperature (K)')
plt.legend(loc='best')

plt.subplot(1,2,2)
plt.hold('on')
names = anames[4:6]
ref = anames[3]
for name in names:
    plt.plot(atms[name].t-atms[ref].t,atms[name].p,label=name)
    plt.plot(atms[name].tconv-atms[ref].t[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(atms[name].tcold-atms[ref].t[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
plt.ylim(yls)
plt.xlim(xls1)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Temperature (K)')
plt.legend(loc='best')

figname = 'img/rce_{}g{}_cld_{}_tdiff.png'.format(
                 radmodel,fstag,rhlabel)
print("Writing figure: {}".format(figname))
plt.savefig(
    bbox_inches='tight', dpi=300, filename=figname)

# %% plot of cloud heating rates
plt.figure(8)
plt.clf()
xls = (-0.5,1.5)
yls = (80, 1000)

plt.subplot(131)
plt.hold('on')
names = anames[1:3]
for name in names:
    plt.plot(auxhr[name].hr,atms[name].p,label=name)
    plt.plot(auxhr[name].hr[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hr[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
names = anames[4:6]
for name in names:
    plt.plot(auxhr[name].hr,atms[name].p,'--',label=name)
    plt.plot(auxhr[name].hr[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hr[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
plt.ylim(yls)
plt.xlim(xls)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Cloud HR (K/d)')
#plt.legend(loc='best')

plt.subplot(132)
plt.hold('on')
names = anames[1:3]
for name in names:
    plt.plot(auxhr[name].hrir,atms[name].p,label=name)
    plt.plot(auxhr[name].hrir[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hrir[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
names = anames[4:6]
for name in names:
    plt.plot(auxhr[name].hrir,atms[name].p,'--',label=name)
    plt.plot(auxhr[name].hrir[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hrir[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
plt.ylim(yls)
plt.xlim(xls)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Cloud IR HR (K/d)')
#plt.legend(loc='best')

plt.subplot(133)
plt.hold('on')
names = anames[1:3]
for name in names:
    plt.plot(auxhr[name].hrsw,atms[name].p,label=name)
    plt.plot(auxhr[name].hrsw[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hrsw[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
names = anames[4:6]
for name in names:
    plt.plot(auxhr[name].hrsw,atms[name].p,'--',label=name)
    plt.plot(auxhr[name].hrsw[atms[name].iconv], atms[name].pconv, 'ks')
    plt.plot(auxhr[name].hrsw[atms[name].icold], atms[name].pcold, 'ko')
    plt.plot(np.zeros(len(atms[name])),atms[name].p, 'k--')
plt.ylim(yls)
plt.xlim(xls)
ax = plt.gca()
ax.set_yscale('log')
ax.invert_yaxis()
plt.title('')
plt.ylabel('Pressure (hPa)')
plt.xlabel('Cloud SW HR (K/d)')
#plt.legend(loc='best')


# required tensoefrlow version: 1.14.0
# conda install -c anaconda tensorflow-gpu==1.14.0

import Utils
# import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import CNN_Siamese
import ResNet_Siamese
import numpy as np
import matplotlib.pyplot as plt
# import cv2
from collections import OrderedDict  #--> for not repeating legends in plot
import umap
import os
from Evaluate_embedding_space import Evaluate_embedding_space
import dataset_characteristics

# import warnings
# warnings.filterwarnings('ignore')

def main():
    #================================ settings:
    train_the_embedding_space = True
    evaluate_the_embedding_space = False
    assert train_the_embedding_space != evaluate_the_embedding_space
    deep_model = "ResNet"  #--> "CNN", "ResNet"
    loss_type = "triplet"   #--> "triplet", "FDA", "contrastive", "FDA_contrastive"
    n_res_blocks = 18  #--> 18, 34, 50, 101, 152
    batch_size = 32
    learning_rate = 1e-5
    margin_in_loss = 0.25
    latent_space_dimension = 100
    feature_space_dimension = 2
    path_save_network_model = ".\\network_model\\" + deep_model + "\\"
    model_dir_ = model_dir(model_name=deep_model, n_res_blocks=n_res_blocks, batch_size=batch_size, learning_rate=learning_rate)
    #================================
    if train_the_embedding_space:
        train_embedding_space(deep_model, n_res_blocks, batch_size, learning_rate, path_save_network_model, model_dir_, feature_space_dimension, latent_space_dimension, margin_in_loss, loss_type)
    if evaluate_the_embedding_space:
        evaluate_embedding_space(path_save_network_model, model_dir_, deep_model, feature_space_dimension, latent_space_dimension, n_res_blocks, margin_in_loss, loss_type)

def evaluate_embedding_space(path_save_network_model, model_dir_, deep_model, feature_space_dimension, latent_space_dimension, n_res_blocks, margin_in_loss, loss_type):
    which_epoch_to_load_NN_model = 45
    path_save_embeddings_of_test_data = ".\\results\\" + deep_model + "\\embedding_test_set\\"
    image_height = dataset_characteristics.get_image_height()
    image_width = dataset_characteristics.get_image_width()
    image_n_channels = dataset_characteristics.get_image_n_channels()
    # path_save_network_model = "./network_model/ResNet/"
    if deep_model == "CNN":
        siamese = CNN_Siamese.CNN_Siamese(loss_type=loss_type, feature_space_dimension=feature_space_dimension, margin_in_loss=margin_in_loss)
    elif deep_model == "ResNet":
        siamese = ResNet_Siamese.ResNet_Siamese(loss_type=loss_type, feature_space_dimension=feature_space_dimension, latent_space_dimension=latent_space_dimension,
                                                n_res_blocks=n_res_blocks, margin_in_loss=margin_in_loss, is_train=True)
    evaluate_ = Evaluate_embedding_space(checkpoint_dir=path_save_network_model+str(which_epoch_to_load_NN_model)+"/", model_dir_=model_dir_)
    (X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
    # X_train = X_train[:2000, :, :]
    # X_train = X_train.reshape((X_train.shape[0], image_height, image_width, image_n_channels))
    X_test = X_test.reshape((X_test.shape[0], image_height, image_width, image_n_channels))
    # evaluate_.embed_the_data(X=X_train, labels=y_train, siamese=siamese, path_save_embeddings_of_test_data=path_save_embeddings_of_test_data)
    embedding, labels = evaluate_.embed_the_data(X=X_test, labels=y_test, siamese=siamese, path_save_embeddings_of_test_data=path_save_embeddings_of_test_data)
    evaluate_.classify_with_1NN(embedding, labels, path_to_save=path_save_embeddings_of_test_data + "KNN/")

def train_embedding_space(deep_model, n_res_blocks, batch_size, learning_rate, path_save_network_model, model_dir_, feature_space_dimension, latent_space_dimension, margin_in_loss, loss_type):
    #================================ settings:
    save_plot_embedding_space = True
    save_points_in_embedding_space = True
    load_saved_network_model = False
    which_epoch_to_load_NN_model = 0
    num_epoch = 51
    save_network_model_every_how_many_epochs = 2
    save_embedding_every_how_many_epochs = 2
    # STEPS_PER_EPOCH_TRAIN = 704
    # STEPS_PER_EPOCH_TRAIN = 16
    STEPS_PER_EPOCH_TRAIN = 312  #--> 10000/32 (n_samples/batch_size)
    n_samples_plot = None   #--> if None, plot all
    image_height = dataset_characteristics.get_image_height()
    image_width = dataset_characteristics.get_image_width()
    image_n_channels = dataset_characteristics.get_image_n_channels()
    # path_tfrecords_train = 'C:\\Users\\bghojogh\\Desktop\\My_PhD\\PhD_projects\\Pathology\\dataset\\TCGA_triplets\\tfrecord\\triplets.tfrecords'
    # path_tfrecords_train = 'C:\\Users\\bghojogh\\Desktop\\My_PhD\\PhD_projects\\Fisher_loss\\codes\\4_make_triplets\\2_MNIST\\triplets\\MNIST_1024_triplets\\tfrecord\\triplets.tfrecords'
    # path_tfrecords_train = 'C:\\Users\\bghojogh\\Desktop\\My_PhD\\PhD_projects\\Fisher_loss\\codes\\4_make_triplets\\2_MNIST\\triplets\\MNIST_500_triplets\\tfrecord\\triplets.tfrecords'
    path_tfrecords_train = '.\\triplets.tfrecords'
    path_save_embedding_space = ".\\results\\" + deep_model + "\\embedding_train_set\\"
    path_save_loss = ".\\loss_saved\\"
    #================================

    train_dataset = tf.data.TFRecordDataset([path_tfrecords_train])
    train_dataset = train_dataset.map(Utils.parse_function)
    train_dataset = train_dataset.map(Utils.normalize_triplets)

    num_repeat = None
    train_dataset = train_dataset.repeat(num_repeat)
    train_dataset = train_dataset.shuffle(buffer_size=1024)
    train_dataset = train_dataset.batch(batch_size)
    handle = tf.placeholder(tf.string, shape=[])
    iterator = tf.data.Iterator.from_string_handle(handle, train_dataset.output_types,
                                                             train_dataset.output_shapes)

    next_element = iterator.get_next()
    # training_iterator = train_dataset.make_initializable_iterator()
    training_iterator = tf.data.make_initializable_iterator(train_dataset)

    # Siamese:
    if deep_model == "CNN":
        siamese = CNN_Siamese.CNN_Siamese(loss_type=loss_type, feature_space_dimension=feature_space_dimension, margin_in_loss=margin_in_loss)
    elif deep_model == "ResNet":
        siamese = ResNet_Siamese.ResNet_Siamese(loss_type=loss_type, feature_space_dimension=feature_space_dimension, latent_space_dimension=latent_space_dimension,
                                                n_res_blocks=n_res_blocks, margin_in_loss=margin_in_loss, is_train=True)
    # train_step = tf.train.GradientDescentOptimizer(learning_rate=0.1).minimize(siamese.loss)
    train_step = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(siamese.loss)
    # tf.initialize_all_variables().run()

    saver_ = tf.train.Saver(max_to_keep=None)  # https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/Saver

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        training_handle = sess.run(training_iterator.string_handle())
        sess.run(training_iterator.initializer)

        if load_saved_network_model:
            succesful_load, latest_epoch = load_network_model(saver_=saver_, session_=sess, checkpoint_dir=path_save_network_model+str(which_epoch_to_load_NN_model)+"/",
                                                                model_dir_=model_dir_, model_name=deep_model)
            assert (succesful_load == True)
            loss_average_of_epochs = np.load(path_save_loss + "loss.npy")
            loss_average_of_epochs = loss_average_of_epochs[:latest_epoch+1]
            loss_average_of_epochs = list(loss_average_of_epochs)
        else:
            latest_epoch = -1
            loss_average_of_epochs = []

        for epoch in range(latest_epoch+1, num_epoch):
            losses_in_epoch = []
            print("============= epoch: " + str(epoch) + "/" + str(num_epoch-1))
            embeddings_in_epoch = np.zeros((STEPS_PER_EPOCH_TRAIN * batch_size * 3, feature_space_dimension))
            labels_in_epoch = np.zeros((STEPS_PER_EPOCH_TRAIN * batch_size * 3,))
            latent_space_dimension = 100  #--> see the file ResNet_siamese.py
            embeddings_in_epoch_secondToLast = np.zeros((STEPS_PER_EPOCH_TRAIN * batch_size * 3, latent_space_dimension))
            for i in range(STEPS_PER_EPOCH_TRAIN):
                image_anchor, image_neighbor, image_distant, label_anchor, label_neighbor, label_distant = sess.run(next_element,
                                                                       feed_dict={handle: training_handle})

                image_anchor = image_anchor.reshape((batch_size, image_height, image_width, image_n_channels))
                image_neighbor = image_neighbor.reshape((batch_size, image_height, image_width, image_n_channels))
                image_distant = image_distant.reshape((batch_size, image_height, image_width, image_n_channels))

                _, loss_v, embedding1, embedding2, embedding3, \
                embedding1_secondToLast, embedding2_secondToLast, embedding3_secondToLast = sess.run([train_step, siamese.loss, siamese.o1, siamese.o2, siamese.o3,
                                                                                siamese.o1_secondToLast, siamese.o2_secondToLast, siamese.o3_secondToLast], feed_dict={
                                                                                siamese.x1: image_anchor,
                                                                                siamese.x2: image_neighbor,
                                                                                siamese.x3: image_distant})

                embeddings_in_epoch[ ((i*3*batch_size)+(0*batch_size)) : ((i*3*batch_size)+(1*batch_size)), : ] = embedding1
                embeddings_in_epoch[ ((i*3*batch_size)+(1*batch_size)) : ((i*3*batch_size)+(2*batch_size)), : ] = embedding2
                embeddings_in_epoch[ ((i*3*batch_size)+(2*batch_size)) : ((i*3*batch_size)+(3*batch_size)), : ] = embedding3

                labels_in_epoch[ ((i*3*batch_size)+(0*batch_size)) : ((i*3*batch_size)+(1*batch_size)) ] = label_anchor
                labels_in_epoch[ ((i*3*batch_size)+(1*batch_size)) : ((i*3*batch_size)+(2*batch_size)) ] = label_neighbor
                labels_in_epoch[ ((i*3*batch_size)+(2*batch_size)) : ((i*3*batch_size)+(3*batch_size)) ] = label_distant

                embeddings_in_epoch_secondToLast[ ((i*3*batch_size)+(0*batch_size)) : ((i*3*batch_size)+(1*batch_size)), : ] = embedding1_secondToLast
                embeddings_in_epoch_secondToLast[ ((i*3*batch_size)+(1*batch_size)) : ((i*3*batch_size)+(2*batch_size)), : ] = embedding2_secondToLast
                embeddings_in_epoch_secondToLast[ ((i*3*batch_size)+(2*batch_size)) : ((i*3*batch_size)+(3*batch_size)), : ] = embedding3_secondToLast

                losses_in_epoch.extend([loss_v])

            # report average loss of epoch:
            loss_average_of_epochs.append(np.average(np.asarray(losses_in_epoch)))
            print("Average loss of epoch " + str(epoch) + ": " + str(loss_average_of_epochs[-1]))
            if not os.path.exists(path_save_loss):
                os.makedirs(path_save_loss)
            np.save(path_save_loss + "loss.npy", np.asarray(loss_average_of_epochs))

            # plot the embedding space:
            if (epoch % save_embedding_every_how_many_epochs == 0):
                if save_points_in_embedding_space:
                    if not os.path.exists(path_save_embedding_space+"numpy\\"):
                        os.makedirs(path_save_embedding_space+"numpy\\")
                    np.save(path_save_embedding_space+"numpy\\embeddings_in_epoch_" + str(epoch) + ".npy", embeddings_in_epoch)
                    np.save(path_save_embedding_space+"numpy\\labels_in_epoch_" + str(epoch) + ".npy", labels_in_epoch)
                    # np.save(path_save_embedding_space+"numpy\\embeddings_in_epoch_secondToLast_" + str(epoch) + ".npy", embeddings_in_epoch_secondToLast)
                if save_plot_embedding_space:
                    print("saving the plot of embedding space....")
                    plt.figure(200)
                    # fig.clf()
                    _, indices_to_plot = plot_embedding_of_points(embeddings_in_epoch, labels_in_epoch, n_samples_plot)
                    if not os.path.exists(path_save_embedding_space+"plots\\"):
                        os.makedirs(path_save_embedding_space+"plots\\")
                    plt.savefig(path_save_embedding_space+"plots\\" + 'epoch' + str(epoch) + '_step' + str(i) + '.png')
                    plt.clf()
                    plt.close()
                    # if not os.path.exists(path_save_embedding_space+"plots_secondToLast\\"):
                    #     os.makedirs(path_save_embedding_space+"plots_secondToLast\\")
                    # plot_embedding_of_points_secondToLast(embeddings_in_epoch_secondToLast, labels_in_epoch, indices_to_plot)
                    # plt.savefig(path_save_embedding_space+"plots_secondToLast\\" + 'epoch' + str(epoch) + '_step' + str(i) + '.png')
                    # plt.clf()
                    # plt.close()

            # save the network model:
            if (epoch % save_network_model_every_how_many_epochs == 0):
                # save_network_model(saver_=saver_, session_=sess, checkpoint_dir=path_save_network_model, step=epoch, model_name=deep_model, model_dir_=model_dir_)
                save_network_model(saver_=saver_, session_=sess, checkpoint_dir=path_save_network_model+str(epoch)+"/", step=epoch, model_name=deep_model, model_dir_=model_dir_)
                print("Model saved in path: %s" % path_save_network_model)

def plot_embedding_of_points(embedding, labels, n_samples_plot=None):
    n_samples = embedding.shape[0]
    if n_samples_plot != None:
        indices_to_plot = np.random.choice(range(n_samples), min(n_samples_plot, n_samples), replace=False)
    else:
        indices_to_plot = np.random.choice(range(n_samples), n_samples, replace=False)
    embedding_sampled = embedding[indices_to_plot, :]
    if embedding.shape[1] == 2:
        pass
    else:
        embedding_sampled = umap.UMAP(n_neighbors=500).fit_transform(embedding_sampled)
    n_points = embedding.shape[0]
    # n_points_sampled = embedding_sampled.shape[0]
    labels_sampled = labels[indices_to_plot]
    _, ax = plt.subplots(1, figsize=(14, 10))
    classes = dataset_characteristics.get_class_names()
    n_classes = len(classes)
    plt.scatter(embedding_sampled[:, 0], embedding_sampled[:, 1], s=10, c=labels_sampled, cmap='Spectral', alpha=1.0)
    # plt.setp(ax, xticks=[], yticks=[])
    cbar = plt.colorbar(boundaries=np.arange(n_classes+1)-0.5)
    cbar.set_ticks(np.arange(n_classes))
    cbar.set_ticklabels(classes)
    return plt, indices_to_plot

def plot_embedding_of_points_secondToLast(embedding_secondToLast, labels, indices_to_plot):
    embedding_sampled = embedding_secondToLast[indices_to_plot, :]
    if embedding_secondToLast.shape[1] == 2:
        pass
    else:
        embedding_sampled = umap.UMAP(n_neighbors=500).fit_transform(embedding_sampled)
    n_points = embedding_secondToLast.shape[0]
    # n_points_sampled = embedding_sampled.shape[0]
    labels_sampled = labels[indices_to_plot]
    _, ax = plt.subplots(1, figsize=(14, 10))
    classes = dataset_characteristics.get_class_names()
    n_classes = len(classes)
    plt.scatter(embedding_sampled[:, 0], embedding_sampled[:, 1], s=10, c=labels_sampled, cmap='Spectral', alpha=1.0)
    # plt.setp(ax, xticks=[], yticks=[])
    cbar = plt.colorbar(boundaries=np.arange(n_classes+1)-0.5)
    cbar.set_ticks(np.arange(n_classes))
    cbar.set_ticklabels(classes)
    return plt

def save_network_model(saver_, session_, checkpoint_dir, step, model_name, model_dir_):
    # https://stackoverflow.com/questions/33759623/tensorflow-how-to-save-restore-a-model
    # https://github.com/taki0112/ResNet-Tensorflow/blob/master/ResNet.py
    checkpoint_dir = os.path.join(checkpoint_dir, model_dir_)
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    saver_.save(session_, os.path.join(checkpoint_dir, model_name+'.model'), global_step=step)

def load_network_model(saver_, session_, checkpoint_dir, model_dir_, model_name):
    # https://stackoverflow.com/questions/33759623/tensorflow-how-to-save-restore-a-model
    print(" [*] Reading checkpoints...")
    checkpoint_dir = os.path.join(checkpoint_dir, model_dir_)
    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
    if ckpt and ckpt.model_checkpoint_path:
        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
        saver_.restore(session_, os.path.join(checkpoint_dir, ckpt_name))
        print(" [*] Success to read {}".format(ckpt_name))
        latest_epoch = int(ckpt_name.split("-")[-1])
        return True, latest_epoch
    else:
        print(" [*] Failed to find a checkpoint")
        return False, 0

def model_dir(model_name, n_res_blocks, batch_size, learning_rate):
    return "{}_{}_{}_{}".format(model_name, n_res_blocks, batch_size, learning_rate)

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
"""
Created on Thu Jun 16 13:24:45 2016

@author: Ent00002
"""

# delayed runs, comment the 2 lines below when unused
# import time
# time.sleep(100)

#%% Import libraries

import numpy as np
import scipy.io as sio
import calendar
import datetime
from getconstants import getconstants
from timeit import default_timer as timer
import os

#%%BEGIN OF INPUT1 (FILL THIS IN)
years		= np.arange(1979,2016)	# Note that this can be forward as tracking was already finished, whereas the masterscript must be in backward order
yearpart	= np.arange(0,366)	# for a full (leap)year fill in np.arange(0,366)
daily		= 0			# 1 for writing out daily data, 0 for only monthly data
timetracking	= 1			# 0 for not tracking time and 1 for tracking time

# Manage the extent of your dataset (FILL THIS IN)
# Define the latitude and longitude cell numbers to consider and corresponding lakes that should be considered part of the land
latnrs = np.arange(7,114)
lonnrs = np.arange(0,240)

# the lake numbers below belong to the ERA-Interim data on 1.5 degree starting at Northern latitude 79.5 and longitude -180
lake_mask_1 = np.array([9,9,9,12,12,21,21,22,22,23,24,25,23,23,25,25,53,54,61,23,24,23,24,25,27,22,23,24,25,26,27,28,22,25,26,27,28,23,23,12,18])
lake_mask_2 = np.array([120+19,120+40,120+41,120+43,120+44,120+61,120+62,120+62,120+63,120+62,120+62,120+62,120+65,120+66,120+65,120+66,142-120,142-120,143-120,152-120,152-120,153-120,153-120,153-120,153-120,154-120,154-120,154-120,154-120,154-120,154-120,154-120,155-120,155-120,155-120,155-120,155-120,159-120,160-120,144-120,120+55])
lake_mask = np.transpose(np.vstack((lake_mask_1,lake_mask_2))) #recreate the arrays of the matlab model

# obtain the constants
invariant_data = '/home/users/lguo/CSSP/WAM2layersPython-master/download_scripts/invariants.nc'	#invariants
latitude,longitude,lsm,g,density_water,timestep,A_gridcell,L_N_gridcell,L_S_gridcell,L_EW_gridcell,gridcell = getconstants(latnrs,lonnrs,lake_mask,invariant_data)

interdata_folder = r'/home/users/lguo/CSSP/WAM2layersPython-master/interdata'
output_folder = r'/home/users/lguo/CSSP/WAM2layersPython-master/output'
sub_interdata_folder = os.path.join(interdata_folder, 'cn1_backward')

#END OF INPUT

#%% Datapaths (FILL THIS IN)

def data_path(y,a,years,timetracking):
    load_Sa_track = os.path.join(sub_interdata_folder, str(y) + '-' + str(a) + 'Sa_track.mat')
    load_Sa_time = os.path.join(sub_interdata_folder, str(y) + '-' + str(a) + 'Sa_time.mat')
    load_fluxes_and_storages = os.path.join(interdata_folder, str(y) + '-' + str(a) + 'fluxes_storages.mat')

    save_path = os.path.join(output_folder, 'E_track_cn1_full' + str(years[0]) + '-' + str(years[-1]) + '-timetracking' + str(timetracking) + '.mat')
    save_path_daily = os.path.join(output_folder, 'E_track_cn1_daily_full' + str(y) + '-timetracking' + str(timetracking) + '.mat')

    return load_Sa_track,load_Sa_time,load_fluxes_and_storages,save_path,save_path_daily

#%% Runtime & Results
start1 = timer()
startyear = years[0]

E_per_year_per_month		= np.zeros((len(years),12,len(latitude),len(longitude)))
E_track_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
P_per_year_per_month		= np.zeros((len(years),12,len(latitude),len(longitude)))
Sa_track_down_per_year_per_month= np.zeros((len(years),12,len(latitude),len(longitude)))
Sa_track_top_per_year_per_month = np.zeros((len(years),12,len(latitude),len(longitude)))
W_down_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
W_top_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
north_loss_per_year_per_month	= np.zeros((len(years),12,1,len(longitude)))
south_loss_per_year_per_month	= np.zeros((len(years),12,1,len(longitude)))
down_to_top_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
top_to_down_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
water_lost_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))

if timetracking == 1:
    Sa_time_down_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
    Sa_time_top_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
    E_time_per_year_per_month		= np.zeros((len(years),12,len(latitude),len(longitude)))
    Sa_dist_down_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
    Sa_dist_top_per_year_per_month	= np.zeros((len(years),12,len(latitude),len(longitude)))
    E_dist_per_year_per_month		= np.zeros((len(years),12,len(latitude),len(longitude)))

for i in range(len(years)):
    y = years[i]
    ly = int(calendar.isleap(y))
    final_time = 364+ly

    E_per_day			= np.zeros((365+ly,len(latitude),len(longitude)))
    E_track_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    P_per_day			= np.zeros((365+ly,len(latitude),len(longitude)))
    Sa_track_down_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
    Sa_track_top_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
    W_down_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    W_top_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    north_loss_per_day		= np.zeros((365+ly,1,len(longitude)))
    south_loss_per_day		= np.zeros((365+ly,1,len(longitude)))
    down_to_top_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    top_to_down_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    water_lost_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
    if timetracking == 1:
        Sa_time_down_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
        Sa_time_top_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
        E_time_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))
        Sa_dist_down_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
        Sa_dist_top_per_day	= np.zeros((365+ly,len(latitude),len(longitude)))
        E_dist_per_day		= np.zeros((365+ly,len(latitude),len(longitude)))

    for j in range(len(yearpart)):
        start = timer()
        a = yearpart[j]
        datapath = data_path(y,a,years,timetracking)
        if a > final_time: # a = 365 (366th index) and not a leapyear\
            pass
        else:
            # load tracked data [97,107,240]
            loading_ST = sio.loadmat(datapath[0],verify_compressed_data_integrity=False)
            Sa_track_top	= loading_ST['Sa_track_top']
            Sa_track_down	= loading_ST['Sa_track_down']
            north_loss		= loading_ST['north_loss']
            south_loss		= loading_ST['south_loss']
            down_to_top		= loading_ST['down_to_top']
            top_to_down		= loading_ST['top_to_down']
            water_lost		= loading_ST['water_lost']
            Sa_track		= Sa_track_top + Sa_track_down
            if timetracking == 1:
                loading_STT = sio.loadmat(datapath[1],verify_compressed_data_integrity=False)
                Sa_time_top	= loading_STT['Sa_time_top']
                Sa_time_down	= loading_STT['Sa_time_down']
                Sa_dist_top	= loading_STT['Sa_dist_top']
                Sa_dist_down	= loading_STT['Sa_dist_down']

            # load the total moisture data [96/97,107,240]
            loading_FS = sio.loadmat(datapath[2],verify_compressed_data_integrity=False)
            Fa_E_top	= loading_FS['Fa_E_top']
            Fa_N_top	= loading_FS['Fa_N_top']
            Fa_E_down	= loading_FS['Fa_E_down']
            Fa_N_down	= loading_FS['Fa_N_down']
            Fa_Vert	= loading_FS['Fa_Vert']
            E		= loading_FS['E']
            P		= loading_FS['P']
            W_top	= loading_FS['W_top']
            W_down	= loading_FS['W_down']
            W = W_top + W_down

            # compute tracked evaporation [96,107,240]
	    # LG: E_track is calculated via the bottom layer Sa.
            E_track = E[:,:,:] * (Sa_track_down[1:,:,:] / W_down[1:,:,:])

            # save per day
            E_per_day[a,:,:]			= np.sum(E, axis =0)
            E_track_per_day[a,:,:]		= np.sum(E_track, axis =0)
            P_per_day[a,:,:]			= np.sum(P, axis =0)
            Sa_track_down_per_day[a,:,:]	= np.mean(Sa_track_down[1:,:,:], axis =0)
            Sa_track_top_per_day[a,:,:]		= np.mean(Sa_track_top[1:,:,:], axis =0)
            W_down_per_day[a,:,:]		= np.mean(W_down[1:,:,:], axis =0)
            W_top_per_day[a,:,:]		= np.mean(W_top[1:,:,:], axis =0)

            north_loss_per_day[a,:,:]		= np.sum(north_loss, axis =0)
            south_loss_per_day[a,:,:]		= np.sum(south_loss, axis =0)
            down_to_top_per_day[a,:,:]		= np.sum(down_to_top, axis =0)
            top_to_down_per_day[a,:,:]		= np.sum(top_to_down, axis =0)
            water_lost_per_day[a,:,:]		= np.sum(water_lost, axis =0)

            if timetracking == 1:
                # compute tracked evaporation time [96,107,240]
		# LG: E_time is mean of two concecutive timesteps
                E_time = 0.5 * ( Sa_time_down[:-1,:,:] + Sa_time_down[1:,:,:] )					# seconds
                # save per day
                Sa_time_down_per_day[a,:,:]	= np.mean(Sa_time_down[:-1,:,:], axis=0)			# seconds
                Sa_time_top_per_day[a,:,:]	= np.mean(Sa_time_top[:-1,:,:], axis=0)				# seconds
                E_time_per_day[a,:,:]		= np.sum((E_time * E_track), axis = 0) / E_track_per_day[a,:,:]	# seconds
                # remove nans
                where_are_NaNs = np.isnan(E_time_per_day)
                E_time_per_day[where_are_NaNs] = 0

                # compute tracked evaporation distance [96,107,240]
                E_dist = 0.5 * ( Sa_dist_down[:-1,:,:] + Sa_dist_down[1:,:,:] )					# m
                # save per day
                Sa_dist_down_per_day[a,:,:]	= np.mean(Sa_dist_down[:-1,:,:], axis=0)			# m
                Sa_dist_top_per_day[a,:,:]	= np.mean(Sa_dist_top[:-1,:,:], axis=0)				# m
                E_dist_per_day[a,:,:]		= np.sum((E_dist * E_track), axis = 0) / E_track_per_day[a,:,:]	# m
                # remove nans
                where_are_NaNs = np.isnan(E_dist_per_day)
                E_dist_per_day[where_are_NaNs] = 0

        end = timer()
        print 'Runtime output for day ' + str(a+1) + ' in year ' + str(y) + ' is',(end - start),' seconds.'

    if daily == 1:
        if timetracking == 0: # create dummy values
            Sa_time_down_per_day	= 0
            Sa_time_top_per_day		= 0
            E_time_per_day		= 0

        sio.savemat(datapath[4],
                    {'E_per_day':E_per_day,'E_track_per_day':E_track_per_day,'P_per_day':P_per_day,
                     'Sa_track_down_per_day':Sa_track_down_per_day,'Sa_track_top_per_day':Sa_track_top_per_day,
                     'Sa_time_down_per_day':Sa_time_down_per_day,'Sa_time_top_per_day':Sa_time_top_per_day,
                     'Sa_dist_down_per_day':Sa_dist_down_per_day,'Sa_dist_top_per_day':Sa_dist_top_per_day,
                     'W_down_per_day':W_down_per_day,'W_top_per_day':W_top_per_day,
                     'E_time_per_day':E_time_per_day,'E_dist_per_day':E_dist_per_day},do_compression=True)

    # values per month
    for m in range(12):
        first_day = int(datetime.date(y,m+1,1).strftime("%j"))
        last_day = int(datetime.date(y,m+1,calendar.monthrange(y,m+1)[1]).strftime("%j"))
        days = np.arange(first_day,last_day+1)-1						# -1 because Python is zero-based

        E_per_year_per_month[y-startyear,m,:,:]			= (np.squeeze(np.sum(E_per_day[days,:,:], axis = 0)))
        E_track_per_year_per_month[y-startyear,m,:,:]		= (np.squeeze(np.sum(E_track_per_day[days,:,:], axis = 0)))
        P_per_year_per_month[y-startyear,m,:,:]			= (np.squeeze(np.sum(P_per_day[days,:,:], axis = 0)))
        Sa_track_down_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_track_down_per_day[days,:,:], axis = 0)))
        Sa_track_top_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_track_top_per_day[days,:,:], axis = 0)))
        W_down_per_year_per_month[y-startyear,m,:,:]		= (np.squeeze(np.mean(W_down_per_day[days,:,:], axis = 0)))
        W_top_per_year_per_month[y-startyear,m,:,:]		= (np.squeeze(np.mean(W_top_per_day[days,:,:], axis = 0)))
        north_loss_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum(north_loss_per_day[days,:,:], axis = 0)))
        south_loss_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum(south_loss_per_day[days,:,:], axis = 0)))
        down_to_top_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum(down_to_top_per_day[days,:,:], axis = 0)))
        top_to_down_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum(top_to_down_per_day[days,:,:], axis = 0)))
        water_lost_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum(water_lost_per_day[days,:,:], axis = 0)))

        if timetracking == 1:
            # tracked time
            Sa_time_down_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_time_down_per_day[days,:,:], axis = 0)))
            Sa_time_top_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_time_top_per_day[days,:,:], axis =0)))
            E_time_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum( E_time_per_day[days,:,:]
                * E_track_per_day[days,:,:], axis=0)) / np.squeeze(E_track_per_year_per_month[y-startyear,m,:,:]))

            # remove nans
            where_are_NaNs = np.isnan(E_time_per_year_per_month)
            E_time_per_year_per_month[where_are_NaNs] = 0

            # tracked distance
            Sa_dist_down_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_dist_down_per_day[days,:,:], axis = 0)))
            Sa_dist_top_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.mean(Sa_dist_top_per_day[days,:,:], axis =0)))
            E_dist_per_year_per_month[y-startyear,m,:,:]	= (np.squeeze(np.sum( E_dist_per_day[days,:,:]
                * E_track_per_day[days,:,:], axis=0)) / np.squeeze(E_track_per_year_per_month[y-startyear,m,:,:]))

            # remove nans
            where_are_NaNs = np.isnan(E_dist_per_year_per_month)
            E_dist_per_year_per_month[where_are_NaNs] = 0

        elif timetracking == 0:
            Sa_time_down_per_year_per_month	= 0
            Sa_time_top_per_year_per_month	= 0
            E_time_per_year_per_month		= 0
            Sa_dist_down_per_year_per_month	= 0
            Sa_dist_top_per_year_per_month	= 0
            E_dist_per_year_per_month		= 0

# save monthly data
sio.savemat(datapath[3],
           {'E_per_year_per_month':E_per_year_per_month,'E_track_per_year_per_month':E_track_per_year_per_month,'P_per_year_per_month':P_per_year_per_month,
            'Sa_track_down_per_year_per_month':Sa_track_down_per_year_per_month,'Sa_track_top_per_year_per_month':Sa_track_top_per_year_per_month,
            'Sa_time_down_per_year_per_month':Sa_time_down_per_year_per_month,'Sa_time_top_per_year_per_month':Sa_time_top_per_year_per_month,
            'E_time_per_year_per_month':E_time_per_year_per_month, 'Sa_dist_down_per_year_per_month':Sa_dist_down_per_year_per_month,'Sa_dist_top_per_year_per_month':Sa_dist_top_per_year_per_month, 'E_dist_per_year_per_month':E_dist_per_year_per_month, 'W_down_per_year_per_month':W_down_per_year_per_month,'W_top_per_year_per_month':W_top_per_year_per_month,
            'north_loss_per_year_per_month':north_loss_per_year_per_month,'south_loss_per_year_per_month':south_loss_per_year_per_month,
            'down_to_top_per_year_per_month':down_to_top_per_year_per_month,'top_to_down_per_year_per_month':top_to_down_per_year_per_month,
            'water_lost_per_year_per_month':water_lost_per_year_per_month})

end1 = timer()
print 'The total runtime of Con_E_Recyc_Output is',(end1-start1),' seconds.'


#!/usr/bin/env python
"""
Program for the regression of mlp about paramagnetic FCC Fe
"""

import argparse
import copy

# import time
# import tqdm
import random

import numpy as np
from mlptools.common.fileio import InputParams
from mlptools.common.structure import Structure
from mlptools.mlpgen.model import Terms
from mlptools.mlpgen.myIO import ReadFeatureParams, read_regression_params
from mlptools.mlpgen.regression import PotEstimation

def rearange_L(array, index_array):
    """Move designated columns to the head of array.

    Args:
        array (ndarray): input array(2D)
        index_array (list): list of index by which columns are designated

    Returns:
        ndarray: changed array
    """
    rest = np.delete(array, index_array, 1)
    return np.hstack((array[:, index_array], rest))

class VirtualDataInput:
    """Generate a virtual DataInput from normal DataInput"""

    def __init__(self, di):
        self.vdi = copy.deepcopy(di)
        self.vdi.n_type = 2

    def get_data_input(self):
        """Return a newly generated DataInput.

        Returns:
            DataInput: virtual DataInput
        """
        return self.vdi

class MagneticStructuralFeatures:
    """Data structure including magnetic structural features"""

    def __init__(self, tr, spin_array, vdi):
        st_set_all_train = self.get_virtual_structures(tr.train, spin_array)
        n_st_dataset = [len(data.st_set) for data in tr.train]
        term = Terms(st_set_all_train, vdi, n_st_dataset, vdi.train_force)
        self.train_x = np.hstack((tr.train_x, term.get_x()))
        st_set_all_test = self.get_virtual_structures(tr.test, spin_array)
        n_st_dataset = [len(data.st_set) for data in tr.test]
        force_dataset = [vdi.wforce for v in vdi.test_names]
        term = Terms(st_set_all_test, vdi, n_st_dataset, force_dataset)
        self.test_x = np.hstack((tr.test_x, term.get_x()))

    def get_x(self):
        """Return the X matrices for regression

        Returns:
            ndarray: X matrices needed for training and test
        """
        return self.train_x, self.test_x

    def get_virtual_structures(self, dataset, spin_array):
        """Generate virtual structures from dataset based on spin_array and return the
           list of them

        Args:
            dataset (dr_array): array of the instances, DataRegression
            spin_array (ndarray): which spin each atom has

        Returns:
            list: list of rewrited structures
        """
        index_array = np.nonzero(spin_array == 1)[0]
        n_atom_1 = len(index_array)
        n_atom_2 = sum(dataset[0].st_set[0].n_atoms) - len(index_array)
        n_atoms = [n_atom_1, n_atom_2]
        specie1 = ["A" for i in range(n_atom_1)]
        specie2 = ["B" for i in range(n_atom_2)]
        elements = specie1.extend(specie2)
        type1 = [0 for i in range(n_atom_1)]
        type2 = [1 for i in range(n_atom_2)]
        types = type1.extend(type2)
        st_list = [
            Structure(
                st.axis,
                rearange_L(st.positions, index_array),
                n_atoms,
                elements,
                types=types,
                comment=st.comment,
            )
            for data in dataset
            for st in data.st_set
        ]
        return st_list

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--infile",
        type=str,
        required=True,
        help="Input file name. Training is performed from vasprun files.",
    )
    parser.add_argument(
        "-p",
        "--pot",
        type=str,
        default="mlp.pkl",
        help="Potential file name for mlptools",
    )
    args = parser.parse_args()

    # prepare temporary spin array
    spin_array = np.array(random.choices([-1, 1], k=32))

    p = InputParams(args.infile)
    di = ReadFeatureParams(p).get_params()
    vdi = VirtualDataInput(di).get_data_input()

    # calculate structural features
    tr = PotEstimation(di=di)
    # calculate magnetic structural features
    tr.train_x, tr.test_x = MagneticStructuralFeatures(tr, spin_array, vdi).get_x()
    tr.set_regression_data()

    # start regression
    if args.noreg is False:
        reg_method, alpha_min, alpha_max, n_alpha = read_regression_params(p)
        if reg_method == "ridge" or reg_method == "lasso":
            pot = tr.regularization_reg(
                method=reg_method,
                alpha_min=alpha_min,
                alpha_max=alpha_max,
                n_alpha=n_alpha,
                svd=args.svd,
            )
        elif reg_method == "normal":
            pot = tr.normal_reg()

        pot.save_pot(file_name=args.pot)
        pot.save_pot_for_lammps(file_name=args.lammps)

        print(" --- input parameters ----")
        pot.di.model_e.print()
        print(" --- best model ----")
        if reg_method == "ridge" or reg_method == "lasso":
            print(" alpha = ", tr.best_alpha)

        (
            rmse_train_e,
            rmse_test_e,
            rmse_train_f,
            files_train,
            rmse_test_f,
            rmse_train_s,
            rmse_test_s,
            files_test,
        ) = tr.get_best_rmse()

        print(" -- Prediction Error --")
        for re, rf, rs, f in zip(rmse_train_e, rmse_train_f, rmse_train_s, files_train):
            print(" structures :", f)
            print(" rmse (energy, train) = ", re * 1000, " (meV/atom)")
            if rf is not None:
                print(" rmse (force, train) = ", rf, " (eV/ang)")
                print(" rmse (stress, train) = ", rs, " (GPa)")
        for re, rf, rs, f in zip(rmse_test_e, rmse_test_f, rmse_test_s, files_test):
            print(" structures :", f)
            print(" rmse (energy, test) = ", re * 1000, " (meV/atom)")
            if rf is not None:
                print(" rmse (force, test) = ", rf, " (eV/ang)")
                print(" rmse (stress, test) = ", rs, " (GPa)")


"""
Common solar physics coordinate systems.

This submodule implements various solar physics coordinate frames for use with
the `astropy.coordinates` module.
"""
from contextlib import contextmanager

import numpy as np

import astropy.units as u
from astropy.coordinates import ConvertError, QuantityAttribute
from astropy.coordinates.baseframe import BaseCoordinateFrame, RepresentationMapping
from astropy.coordinates.representation import (
    CartesianDifferential,
    CartesianRepresentation,
    CylindricalRepresentation,
    SphericalDifferential,
    SphericalRepresentation,
    UnitSphericalRepresentation,
)
from astropy.time import Time

from sunpy import log
from sunpy.sun.constants import radius as _RSUN
from sunpy.time.time import _variables_for_parse_time_docstring
from sunpy.util.decorators import add_common_docstring
from sunpy.util.exceptions import warn_user
from .frameattributes import ObserverCoordinateAttribute, TimeFrameAttributeSunPy

_J2000 = Time('J2000.0', scale='tt')

__all__ = ['SunPyBaseCoordinateFrame', 'BaseHeliographic',
           'HeliographicStonyhurst', 'HeliographicCarrington',
           'Heliocentric', 'Helioprojective',
           'HeliocentricEarthEcliptic', 'GeocentricSolarEcliptic',
           'HeliocentricInertial', 'GeocentricEarthEquatorial']

def _frame_parameters():
    """
    Returns formatting dictionary to use with add_common_docstring to populate frame docstrings
    """
    ret = {}

    # Each text block is missing the first indent because it already exists in the frame docstring
    ret['data'] = ("data : `~astropy.coordinates.BaseRepresentation` or ``None``\n"
                   "        A representation object or ``None`` to have no data\n"
                   "        (or use the coordinate component arguments, see below).")
    ret['common'] = (f"obstime : {_variables_for_parse_time_docstring()['parse_time_types']}\n"
                     "        The time of the observation.  This is used to determine the\n"
                     "        position of solar-system bodies (e.g., the Sun and the Earth) as\n"
                     "        needed to define the origin and orientation of the frame.\n"
                     "    representation_type : `~astropy.coordinates.BaseRepresentation`, str, optional\n"
                     "        A representation class or string name of a representation class.\n"
                     "        This may change the valid coordinate component arguments from the\n"
                     "        defaults (see above). For example, passing\n"
                     "        ``representation_type='cartesian'`` will make the frame expect\n"
                     "        Cartesian coordinate component arguments (typically, ``x``, ``y``,\n"
                     "        and ``z``).\n"
                     "    copy : bool, optional\n"
                     "        If `True` (default), make copies of the input coordinate arrays.")
    ret['lonlat'] = ("lon : `~astropy.coordinates.Angle` or `~astropy.units.Quantity`, optional\n"
                     "        The longitude coordinate for this object (``lat`` must also be\n"
                     "        given and ``data`` must be ``None``).\n"
                     "        Not needed if ``data`` is given.\n"
                     "    lat : `~astropy.coordinates.Angle` or `~astropy.units.Quantity`, optional\n"
                     "        The latitude coordinate for this object (``lon`` must also be\n"
                     "        given and ``data`` must be ``None``).\n"
                     "        Not needed if ``data`` is given.")
    ret['radius'] = ("radius : `~astropy.units.Quantity`, optional\n"
                     "        The radial distance coordinate from Sun center for this object.\n"
                     "        Defaults to the radius of the Sun. Not needed if ``data`` is given.")
    ret['distance_sun'] = ("distance : `~astropy.units.Quantity`, optional\n"
                           "        The distance coordinate from Sun center for this object.\n"
                           "        Not needed if ``data`` is given.")
    ret['distance_earth'] = ("distance : `~astropy.units.Quantity`, optional\n"
                             "        The distance coordinate from Earth center for this object.\n"
                             "        Not needed if ``data`` is given.")
    ret['xyz'] = ("x : `~astropy.units.Quantity`, optional\n"
                  "        X-axis coordinate for this object. Not needed if ``data`` is given.\n"
                  "    y : `~astropy.units.Quantity`, optional\n"
                  "        Y-axis coordinate for this object. Not needed if ``data`` is given.\n"
                  "    z : `~astropy.units.Quantity`, optional\n"
                  "        Z-axis coordinate for this object. Not needed if ``data`` is given.")
    ret['observer'] = ("observer : `~sunpy.coordinates.frames.HeliographicStonyhurst`, str\n"
                       "        The location of the observer. If a string is provided,\n"
                       "        it must be a solar system body that can be parsed by\n"
                       "        `~sunpy.coordinates.ephemeris.get_body_heliographic_stonyhurst`\n"
                       "        at the time ``obstime``. Defaults to Earth center.")
    ret['rsun'] = ("rsun : `~astropy.units.Quantity`\n"
                   "        The radius of the Sun in length units. Used to convert a 2D\n"
                   "        coordinate (i.e., no ``radius`` component) to a 3D coordinate by\n"
                   "        assuming that the coordinate is on the surface of the Sun. Defaults\n"
                   "        to the photospheric radius as defined in `sunpy.sun.constants`.")
    ret['equinox'] = (f"equinox : {_variables_for_parse_time_docstring()['parse_time_types']}\n"
                      "        The date for the mean vernal equinox.\n"
                      "        Defaults to the J2000.0 equinox.")

    return ret

class SunPyBaseCoordinateFrame(BaseCoordinateFrame):
    """
    Base class for sunpy coordinate frames.

    This class is not intended to be used directly and has no transformations defined.

    * Defines the frame attribute ``obstime`` for observation time.
    * Defines a default wrap angle of 180 degrees for longitude in spherical coordinates,
      which can be overridden via the class variable ``_wrap_angle``.
    * Inject a nice way of representing the object which the coordinate represents.
    """
    obstime = TimeFrameAttributeSunPy()

    default_representation = SphericalRepresentation
    default_differential = SphericalDifferential

    frame_specific_representation_info = {
        SphericalDifferential: [RepresentationMapping('d_lon', 'd_lon', u.arcsec/u.s),
                                RepresentationMapping('d_lat', 'd_lat', u.arcsec/u.s),
                                RepresentationMapping('d_distance', 'd_distance', u.km/u.s)],
    }

    _wrap_angle = 180*u.deg  # for longitude in spherical coordinates

    def __init__(self, *args, **kwargs):
        self.object_name = None

        # If wrap_longitude=False is passed in, do not impose a specific wrap angle for the frame
        if not kwargs.pop('wrap_longitude', True):
            self._wrap_angle = None

        super().__init__(*args, **kwargs)

        # If obstime is specified, treat the default observer (None) as explicitly set
        if self.obstime is not None and self.is_frame_attr_default('observer'):
            self._attr_names_with_defaults.remove('observer')

        return

    def represent_as(self, base, s='base', in_frame_units=False):
        data = super().represent_as(base, s, in_frame_units=in_frame_units)

        # If a frame wrap angle is set, use that wrap angle for any spherical representations.
        if self._wrap_angle is not None and \
           isinstance(data, (UnitSphericalRepresentation, SphericalRepresentation)):
            data.lon.wrap_angle = self._wrap_angle
        return data

    def __str__(self):
        # We override this here so that when you print a SkyCoord it shows the
        # observer as the string and not the whole massive coordinate.
        if getattr(self, "object_name", None):
            return f"<{self.__class__.__name__} Coordinate for '{self.object_name}'>"
        else:
            return super().__str__()

    @property
    def _is_2d(self):
        return (self._data is not None and self._data.norm().unit is u.one
                and u.allclose(self._data.norm(), 1*u.one))

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

        # TODO: Remove this after the minimum Astropy dependency includes astropy/astropy#12005
        cls._fix_property_docstrings()

    @classmethod
    def _fix_property_docstrings(cls):
        # This class method adds docstrings to properties dynamically created by
        # BaseCoordinateFrame.__init_subclass__().  Accordingly, this method needs to itself be
        # called from SunPyBaseCoordinateFrame.__init_subclass__() to work for our subclasses.
        property_docstrings = {
            'default_representation': "Default representation for position data",
            'default_differential': "Default representation for differential data",
            'frame_specific_representation_info': "Mapping for frame-specific component names",
        }
        for prop, docstring in property_docstrings.items():
            if getattr(cls, prop).__doc__ is None:
                setattr(getattr(cls, prop), '__doc__', docstring)

# TODO: Remove this after the minimum Astropy dependency includes astropy/astropy#12005
SunPyBaseCoordinateFrame._fix_property_docstrings()

class BaseHeliographic(SunPyBaseCoordinateFrame):
    """
    Base class for HeliographicCarrington (HGC) and HeliographicStonyhurst (HGS) frames.

    This class is not intended to be used directly and has no transformations defined.
    """
    frame_specific_representation_info = {
        SphericalRepresentation: [RepresentationMapping('lon', 'lon', u.deg),
                                  RepresentationMapping('lat', 'lat', u.deg),
                                  RepresentationMapping('distance', 'radius', None)],
        SphericalDifferential: [RepresentationMapping('d_lon', 'd_lon', u.arcsec/u.s),
                                RepresentationMapping('d_lat', 'd_lat', u.arcsec/u.s),
                                RepresentationMapping('d_distance', 'd_radius', u.km/u.s)],
    }

    rsun = QuantityAttribute(default=_RSUN, unit=u.km)

    def make_3d(self):
        """
        Returns a fully 3D coordinate based on this coordinate.

        If this coordinate is only 2D (i.e., no ``radius`` component) or is a
        unit vector (i.e., the norm of the coordinate is unity), a new
        coordinate is created that corresponds to the surface of the Sun.
        That is, the 3D coordinate will retain the ``lon`` and ``lat``, and
        ``radius`` will be set to the frame's ``rsun`` frame attribute.

        If this coordinate is already fully 3D, it is directly returned, even
        if it does not lie on the surface of the Sun.

        Returns
        -------
        frame : `~sunpy.coordinates.frames.BaseHeliographic`
            The fully 3D coordinate
        """
        if self._is_2d:
            return self.realize_frame(self._data * self.rsun)

        # The coordinate is already 3D
        return self

@add_common_docstring(**_frame_parameters())
class HeliographicStonyhurst(BaseHeliographic):
    """
    A coordinate or frame in the Stonyhurst Heliographic (HGS) system.

    - The origin is the center of the Sun.
    - The Z-axis (+90 degrees latitude) is aligned with the Sun's north pole.
    - The X-axis (0 degrees longitude and 0 degrees latitude) is aligned with the projection of
      the Sun-Earth line onto the Sun's equatorial plane.

    This system is also know as the Heliocentric Earth Equatorial (HEEQ) system when
    represented using Cartesian components.

    A new instance can be created using the following signatures
    (note that if supplied, ``obstime`` and ``representation_type`` must be
    keyword arguments)::

        HeliographicStonyhurst(lon, lat, obstime=obstime)
        HeliographicStonyhurst(lon, lat, radius, obstime=obstime)
        HeliographicStonyhurst(x, y, z, representation_type='cartesian', obstime=obstime)

    Parameters
    ----------
    {data}
    {lonlat}
    {radius}
    {rsun}
    {common}

    Examples
    --------
    >>> from astropy.coordinates import SkyCoord
    >>> import sunpy.coordinates
    >>> import astropy.units as u
    >>> sc = SkyCoord(1*u.deg, 1*u.deg, 2*u.km,
    ...               frame="heliographic_stonyhurst",
    ...               obstime="2010/01/01T00:00:45")
    >>> sc
    <SkyCoord (HeliographicStonyhurst: obstime=2010-01-01T00:00:45.000, rsun=695700.0 km): (lon, lat, radius) in (deg, deg, km)
        (1., 1., 2.)>
    >>> sc.frame
    <HeliographicStonyhurst Coordinate (obstime=2010-01-01T00:00:45.000, rsun=695700.0 km): (lon, lat, radius) in (deg, deg, km)
        (1., 1., 2.)>
    >>> sc = SkyCoord(HeliographicStonyhurst(-10*u.deg, 2*u.deg))
    >>> sc
    <SkyCoord (HeliographicStonyhurst: obstime=None, rsun=695700.0 km): (lon, lat) in deg
        (-10., 2.)>
    >>> sc = SkyCoord(CartesianRepresentation(0*u.km, 45*u.km, 2*u.km),
    ...               obstime="2011/01/05T00:00:50",
    ...               frame="heliographic_stonyhurst")
    >>> sc
    <SkyCoord (HeliographicStonyhurst: obstime=2011-01-05T00:00:50.000, rsun=695700.0 km): (lon, lat, radius) in (deg, deg, km)
    (90., 2.54480438, 45.04442252)>
    """
    name = "heliographic_stonyhurst"

    def _apply_diffrot(self, duration, rotation_model):
        oldrepr = self.spherical

        from sunpy.physics.differential_rotation import diff_rot
        log.debug(f"Applying {duration} of solar rotation")
        newlon = oldrepr.lon + diff_rot(duration,
                                        oldrepr.lat,
                                        rot_type=rotation_model,
                                        frame_time='sidereal')
        newrepr = SphericalRepresentation(newlon, oldrepr.lat, oldrepr.distance)

        return self.realize_frame(newrepr)

@add_common_docstring(**_frame_parameters())
class HeliographicCarrington(BaseHeliographic):
    """
    A coordinate or frame in the Carrington Heliographic (HGC) system.

    - The origin is the center of the Sun.
    - The Z-axis (+90 degrees latitude) is aligned with the Sun's north pole.
    - The X-axis and Y-axis rotate with a period of 25.38 days.

    This system differs from Stonyhurst Heliographic (HGS) in its definition of longitude.  This
    longitude is an "apparent" longitude because it takes into account the time it takes for light
    to travel from the Sun's surface to the observer (see :ref:`sunpy-coordinates-carrington`).
    Thus, the observer needs to be specified to be able to transform to any other coordinate frame.

    A new instance can be created using the following signatures
    (note that if supplied, ``obstime`` and ``observer`` must be a keyword argument)::

        HeliographicCarrington(lon, lat, obstime=obstime, observer=observer)
        HeliographicCarrington(lon, lat, radius, obstime=obstime, observer=observer)

    If you want to define the location in HGC such that the observer for the coordinate frame is
    the same as that location (e.g., the location of an observatory in its corresponding HGC
    frame), use ``observer='self'``::

        HeliographicCarrington(lon, lat, radius, obstime=obstime, observer='self')

    Parameters
    ----------
    {data}
    {lonlat}
    {radius}
    {observer}
    {rsun}
    {common}

    Examples
    --------
    >>> from astropy.coordinates import SkyCoord
    >>> import sunpy.coordinates
    >>> import astropy.units as u
    >>> sc = SkyCoord(1*u.deg, 2*u.deg, 3*u.km,
    ...               frame="heliographic_carrington",
    ...               observer="earth",
    ...               obstime="2010/01/01T00:00:30")
    >>> sc
    <SkyCoord (HeliographicCarrington: obstime=2010-01-01T00:00:30.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (lon, lat, radius) in (deg, deg, km)
        (1., 2., 3.)>

    >>> sc = SkyCoord([1,2,3]*u.deg, [4,5,6]*u.deg, [5,6,7]*u.km,
    ...               obstime="2010/01/01T00:00:45",
    ...               observer="self",
    ...               frame="heliographic_carrington")
    >>> sc
    <SkyCoord (HeliographicCarrington: obstime=2010-01-01T00:00:45.000, rsun=695700.0 km, observer=self): (lon, lat, radius) in (deg, deg, km)
        [(1., 4., 5.), (2., 5., 6.), (3., 6., 7.)]>

    >>> sc = SkyCoord(CartesianRepresentation(0*u.km, 45*u.km, 2*u.km),
    ...               obstime="2011/01/05T00:00:50",
    ...               frame="heliographic_carrington")
    >>> sc
    <SkyCoord (HeliographicCarrington: obstime=2011-01-05T00:00:50.000, rsun=695700.0 km, observer=None): (lon, lat, radius) in (deg, deg, km)
        (90., 2.54480438, 45.04442252)>
    """
    name = "heliographic_carrington"
    _wrap_angle = 360*u.deg

    observer = ObserverCoordinateAttribute(HeliographicStonyhurst)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if not isinstance(self.observer, BaseCoordinateFrame) and self.observer == 'self' and self._is_2d:
            raise ValueError("Full 3D coordinate (including radius) must be specified "
                             "when observer='self'.")

@add_common_docstring(**_frame_parameters())
class Heliocentric(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Heliocentric system, which is observer-based.

    - The origin is the center of the Sun.
    - The Z-axis is aligned with the Sun-observer line.
    - The Y-axis is aligned with the component of the vector to the Sun's north pole that is
      perpendicular to the Z-axis.

    This frame defaults to a Cartesian component representation, which is known as Heliocentric
    Cartesian (HCC).  This frame can also be represented using cylindrical components, where
    where ``rho`` is the impact parameter and ``psi`` is the position angle.
    ``psi`` is measured relative to the west limb, rather than solar north, so is shifted
    by 90 degrees compared to the convention of the Heliocentric Radial (HCR) system.

    A new instance can be created using the following signatures
    (note that if supplied, ``obstime``, ``observer``, and ``representation_type`` must be
    keyword arguments)::

        Heliocentric(x, y, z, obstime=obstime, observer=observer)
        Heliocentric(rho, psi, z, representation_type='cylindrical', obstime=obstime, observer=observer)

    Parameters
    ----------
    {data}
    {xyz}
    {observer}
    {common}

    Examples
    --------

    >>> from astropy.coordinates import SkyCoord, CartesianRepresentation
    >>> import sunpy.coordinates
    >>> import astropy.units as u

    >>> sc = SkyCoord(CartesianRepresentation(10*u.km, 1*u.km, 2*u.km),
    ...               obstime="2011/01/05T00:00:50", observer="earth", frame="heliocentric")
    >>> sc
    <SkyCoord (Heliocentric: obstime=2011-01-05T00:00:50.000, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (x, y, z) in km
        (10., 1., 2.)>

    >>> sc = SkyCoord([1,2]*u.km, [3,4]*u.m, [5,6]*u.cm,
    ...               obstime="2011/01/01T00:00:54", observer="earth", frame="heliocentric")
    >>> sc
    <SkyCoord (Heliocentric: obstime=2011-01-01T00:00:54.000, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (x, y, z) in (km, m, cm)
        [(1., 3., 5.), (2., 4., 6.)]>

    >>> sc = SkyCoord(CylindricalRepresentation(10*u.km, 60*u.deg, 10*u.km),
    ...               obstime="2011/01/05T00:00:50", observer="earth", frame="heliocentric")
    >>> sc
    <SkyCoord (Heliocentric: obstime=2011-01-05T00:00:50.000, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (x, y, z) in km
        (5., 8.66025404, 10.)>
    """
    default_representation = CartesianRepresentation
    default_differential = CartesianDifferential

    frame_specific_representation_info = {
        CylindricalRepresentation: [RepresentationMapping('phi', 'psi', u.deg)]
    }

    observer = ObserverCoordinateAttribute(HeliographicStonyhurst)

    def represent_as(self, base, s='base', in_frame_units=False):
        data = super().represent_as(base, s, in_frame_units=in_frame_units)

        # For cylindrical representations, wrap the `psi` component (natively `phi`) at 360 deg
        if isinstance(data, CylindricalRepresentation):
            data.phi.wrap_at(360*u.deg, inplace=True)
        return data

@add_common_docstring(**_frame_parameters())
class Helioprojective(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Helioprojective Cartesian (HPC) system, which is observer-based.

    - The origin is the location of the observer.
    - ``Tx`` (aka "theta_x") is the angle relative to the plane containing the Sun-observer line
      and the Sun's rotation axis, with positive values in the direction of the Sun's west limb.
    - ``Ty`` (aka "theta_y") is the angle relative to the Sun's equatorial plane, with positive
      values in the direction of the Sun's north pole.
    - ``distance`` is the Sun-observer distance.

    This system is frequently used in a projective form without ``distance`` specified.  For
    observations looking very close to the center of the Sun, where the small-angle approximation
    is appropriate, ``Tx`` and ``Ty`` can be approximated as Cartesian components.

    A new instance can be created using the following signatures
    (note that if supplied, ``obstime`` and ``observer`` must be keyword arguments)::

        Helioprojective(Tx, Ty, obstime=obstime, observer=observer)
        Helioprojective(Tx, Ty, distance, obstime=obstime, observer=observer)

    Parameters
    ----------
    {data}
    Tx : `~astropy.coordinates.Angle` or `~astropy.units.Quantity`
        The theta_x coordinate for this object. Not needed if ``data`` is given.
    Ty : `~astropy.coordinates.Angle` or `~astropy.units.Quantity`
        The theta_y coordinate for this object. Not needed if ``data`` is given.
    distance : `~astropy.units.Quantity`
        The distance coordinate from the observer for this object.
        Not needed if ``data`` is given.
    {observer}
    {rsun}
    {common}

    Examples
    --------
    >>> from astropy.coordinates import SkyCoord
    >>> import sunpy.coordinates
    >>> import astropy.units as u
    >>> sc = SkyCoord(0*u.deg, 0*u.deg, 5*u.km,
    ...               obstime="2010/01/01T00:00:00", observer="earth", frame="helioprojective")
    >>> sc
    <SkyCoord (Helioprojective: obstime=2010-01-01T00:00:00.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty, distance) in (arcsec, arcsec, km)
        (0., 0., 5.)>
    >>> sc = SkyCoord(0*u.deg, 0*u.deg,
    ...               obstime="2010/01/01T00:00:00", observer="earth", frame="helioprojective")
    >>> sc
    <SkyCoord (Helioprojective: obstime=2010-01-01T00:00:00.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty) in arcsec
        (0., 0.)>
    >>> sc = SkyCoord(CartesianRepresentation(1*u.AU, 1e5*u.km, -2e5*u.km),
    ...               obstime="2011/01/05T00:00:50", observer="earth", frame="helioprojective")
    >>> sc
    <SkyCoord (Helioprojective: obstime=2011-01-05T00:00:50.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty, distance) in (arcsec, arcsec, AU)
        (137.87948623, -275.75878762, 1.00000112)>
    """
    frame_specific_representation_info = {
        SphericalRepresentation: [RepresentationMapping('lon', 'Tx', u.arcsec),
                                  RepresentationMapping('lat', 'Ty', u.arcsec),
                                  RepresentationMapping('distance', 'distance', None)],
        SphericalDifferential: [RepresentationMapping('d_lon', 'd_Tx', u.arcsec/u.s),
                                RepresentationMapping('d_lat', 'd_Ty', u.arcsec/u.s),
                                RepresentationMapping('d_distance', 'd_distance', u.km/u.s)],
        UnitSphericalRepresentation: [RepresentationMapping('lon', 'Tx', u.arcsec),
                                      RepresentationMapping('lat', 'Ty', u.arcsec)],
    }

    rsun = QuantityAttribute(default=_RSUN, unit=u.km)
    observer = ObserverCoordinateAttribute(HeliographicStonyhurst)

    @property
    def angular_radius(self):
        """
        Angular radius of the Sun as seen by the observer.

        The ``rsun`` frame attribute is the radius of the Sun in length units.
        The tangent vector from the observer to the edge of the Sun forms a
        right-angle triangle with the radius of the Sun as the far side and the
        Sun-observer distance as the hypotenuse. Thus, the sine of the angular
        radius of the Sun is ratio of these two distances.
        """
        from sunpy.coordinates.sun import _angular_radius  # avoiding a circular import

        if not isinstance(self.observer, HeliographicStonyhurst):
            if self.observer is None:
                raise ValueError("The observer must be defined, not `None`.")
            raise ValueError("The observer must be fully defined by specifying `obstime`.")
        return _angular_radius(self.rsun, self.observer.radius)

    def make_3d(self):
        """
        This method calculates the third coordinate of the Helioprojective
        frame. It assumes that the coordinate point is on the surface of the Sun.

        If a point in the frame is off limb then NaN will be returned.

        Returns
        -------
        new_frame : `~sunpy.coordinates.frames.Helioprojective`
            A new frame instance with all the attributes of the original but
            now with a third coordinate.
        """
        # Skip if we already are 3D
        if not self._is_2d:
            return self

        if not isinstance(self.observer, BaseCoordinateFrame):
            raise ConvertError("Cannot calculate distance to the Sun "
                               f"for observer '{self.observer}' "
                               "without `obstime` being specified.")

        rep = self.represent_as(UnitSphericalRepresentation)
        lat, lon = rep.lat, rep.lon

        # Check for the use of floats with lower precision than the native Python float
        if not set([lon.dtype.type, lat.dtype.type]).issubset([float, np.float64, np.longdouble]):
            warn_user("The Helioprojective component values appear to be lower "
                      "precision than the native Python float: "
                      f"Tx is {lon.dtype.name}, and Ty is {lat.dtype.name}. "
                      "To minimize precision loss, you may want to cast the values to "
                      "`float` or `numpy.float64` via the NumPy method `.astype()`.")

        # Calculate the distance to the surface of the Sun using the law of cosines
        cos_alpha = np.cos(lat) * np.cos(lon)
        c = self.observer.radius**2 - self.rsun**2
        b = -2 * self.observer.radius * cos_alpha
        # Ignore sqrt of NaNs
        with np.errstate(invalid='ignore'):
            d = ((-1*b) - np.sqrt(b**2 - 4*c)) / 2  # use the "near" solution

        if self._spherical_screen:
            sphere_center = self._spherical_screen['center'].transform_to(self).cartesian
            c = sphere_center.norm()**2 - self._spherical_screen['radius']**2
            b = -2 * sphere_center.dot(rep)
            # Ignore sqrt of NaNs
            with np.errstate(invalid='ignore'):
                dd = ((-1*b) + np.sqrt(b**2 - 4*c)) / 2  # use the "far" solution

            d = np.fmin(d, dd) if self._spherical_screen['only_off_disk'] else dd

        return self.realize_frame(SphericalRepresentation(lon=lon,
                                                          lat=lat,
                                                          distance=d))

    _spherical_screen = None

    @classmethod
    @contextmanager
    def assume_spherical_screen(cls, center, only_off_disk=False):
        """
        Context manager to interpret 2D coordinates as being on the inside of a spherical screen.

        The radius of the screen is the distance between the specified ``center`` and Sun center.
        This ``center`` does not have to be the same as the observer location for the coordinate
        frame.  If they are the same, then this context manager is equivalent to assuming that the
        helioprojective "zeta" component is zero.

        This replaces the default assumption where 2D coordinates are mapped onto the surface of the
        Sun.

        Parameters
        ----------
        center : `~astropy.coordinates.SkyCoord`
            The center of the spherical screen
        only_off_disk : `bool`, optional
            If `True`, apply this assumption only to off-disk coordinates, with on-disk coordinates
            still mapped onto the surface of the Sun.  Defaults to `False`.

        Examples
        --------
        .. minigallery:: sunpy.coordinates.Helioprojective.assume_spherical_screen

        >>> import astropy.units as u
        >>> from sunpy.coordinates import Helioprojective
        >>> h = Helioprojective(range(7)*u.arcsec*319, [0]*7*u.arcsec,
        ...                     observer='earth', obstime='2020-04-08')
        >>> print(h.make_3d())
        <Helioprojective Coordinate (obstime=2020-04-08T00:00:00.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty, distance) in (arcsec, arcsec, AU)
            [(   0., 0., 0.99660825), ( 319., 0., 0.99687244),
             ( 638., 0., 0.99778472), ( 957., 0., 1.00103285),
             (1276., 0.,        nan), (1595., 0.,        nan),
             (1914., 0.,        nan)]>

        >>> with Helioprojective.assume_spherical_screen(h.observer):
        ...     print(h.make_3d())
        <Helioprojective Coordinate (obstime=2020-04-08T00:00:00.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty, distance) in (arcsec, arcsec, AU)
            [(   0., 0., 1.00125872), ( 319., 0., 1.00125872),
             ( 638., 0., 1.00125872), ( 957., 0., 1.00125872),
             (1276., 0., 1.00125872), (1595., 0., 1.00125872),
             (1914., 0., 1.00125872)]>

        >>> with Helioprojective.assume_spherical_screen(h.observer, only_off_disk=True):
        ...     print(h.make_3d())
        <Helioprojective Coordinate (obstime=2020-04-08T00:00:00.000, rsun=695700.0 km, observer=<HeliographicStonyhurst Coordinate for 'earth'>): (Tx, Ty, distance) in (arcsec, arcsec, AU)
            [(   0., 0., 0.99660825), ( 319., 0., 0.99687244),
             ( 638., 0., 0.99778472), ( 957., 0., 1.00103285),
             (1276., 0., 1.00125872), (1595., 0., 1.00125872),
             (1914., 0., 1.00125872)]>
        """
        try:
            old_spherical_screen = cls._spherical_screen  # nominally None

            center_hgs = center.transform_to(HeliographicStonyhurst(obstime=center.obstime))
            cls._spherical_screen = {
                'center': center,
                'radius': center_hgs.radius,
                'only_off_disk': only_off_disk
            }
            yield
        finally:
            cls._spherical_screen = old_spherical_screen

@add_common_docstring(**_frame_parameters())
class HeliocentricEarthEcliptic(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Heliocentric Earth Ecliptic (HEE) system.

    - The origin is the center of the Sun.
    - The X-axis (0 degrees longitude and 0 degrees latitude) is aligned with the Sun-Earth line.
    - The Z-axis (+90 degrees latitude) is aligned with the component perpendicular to the X-axis
      of the mean ecliptic pole at the observation time.

    Parameters
    ----------
    {data}
    {lonlat}
    {distance_sun}
    {common}
    """

@add_common_docstring(**_frame_parameters())
class GeocentricSolarEcliptic(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Geocentric Solar Ecliptic (GSE) system.

    - The origin is the center of the Earth.
    - The X-axis (0 degrees longitude and 0 degrees latitude) is aligned with the Earth-Sun line.
    - The Z-axis (+90 degrees latitude) is aligned with the component perpendicular to the X-axis
      of the mean ecliptic pole at the observation time.

    Parameters
    ----------
    {data}
    {lonlat}
    {distance_earth}
    {common}

    Notes
    -----
    Aberration due to Earth motion is not included.
    """

@add_common_docstring(**_frame_parameters())
class HeliocentricInertial(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Heliocentric Inertial (HCI) system.

    - The origin is the center of the Sun.
    - The Z-axis (+90 degrees latitude) is aligned with the Sun's north pole.
    - The X-axis (0 degrees longitude and 0 degrees latitude) is aligned with the solar ascending
      node on the ecliptic (mean J2000.0).

    Parameters
    ----------
    {data}
    {lonlat}
    {distance_sun}
    {common}

    Notes
    -----
    The solar ascending node on the ecliptic lies on the intersection of the solar equatorial
    plane with the ecliptic plane, not on the intersection of the celestial equatorial plane with
    the ecliptic plane.
    """

@add_common_docstring(**_frame_parameters())
class GeocentricEarthEquatorial(SunPyBaseCoordinateFrame):
    """
    A coordinate or frame in the Geocentric Earth Equatorial (GEI) system.

    - The origin is the center of the Earth.
    - The Z-axis (+90 degrees latitude) is aligned with the Earth's north pole.
    - The X-axis (0 degrees longitude and 0 degrees latitude) is aligned with the mean (not true)
      vernal equinox.

    Parameters
    ----------
    {data}
    {lonlat}
    {distance_earth}
    {equinox}
    {common}

    Notes
    -----
    Aberration due to Earth motion is not included.
    """
    equinox = TimeFrameAttributeSunPy(default=_J2000)

import numpy,import astropy,from astropy