import random
import matplotlib.pyplot as plt
import math
import numpy as np
from scipy.stats import norm
import scipy.stats

def uniforme (a,b):
    x=[]
    for _ in range(1000):
        r = round(random.random(), 4)
        x.append(a+(b-a)*r)
    return x
uni=uniforme(1,3)

def TestChiCuadUni(nums):
    print("Test de bondad Chi Cuadrado para distribucion uniforme")
    observado=[]
    esperado=100
    c=1.2
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9), 2)
    for _ in range (10):
        x =0
        for num in nums:
            if  (c-0.2) <= float(num) <= c : x+=1
        observado.append(x)
        c+=0.2
    x2 = 0
    for obs in observado:
        x2 += (((obs-esperado)**2)/esperado)

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

def exponencial(media):
    x = []
    for _ in range(1000):
        r = random.random()
        x += [-media*(np.log(r))]
    return x
expo=exponencial(5)

def TestChiCuadExp(nums):
    print("Test de bondad Chi Cuadrado para distribucion exponencial:")
    observado=[]
    esperado=[]
    c=0.3
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9), 2)
    for _ in range (9):
        x =0
        for num in nums:
            if  (c-0.3) <= float(num) <= c : x+=1
        observado.append(x)
        esperado.append(1000*((1-(math.e)**(-(1/5)*c))-(1-(math.e)**(-(1/5)*(c-0.3)))))
        c+=0.3
    observado.append(1000-sum(observado))
    esperado.append(1000*((math.e)**(-(1/5)*(c-0.3))))
    x2=0
    for i in range(len(observado)):
        x2 += (((observado[i]-esperado[i])**2)/esperado[i])

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

def gamma(k,a):
    x=[]
    for _ in range(1, 1000):
        tr=1.0
        for _ in range(1,k):
            r = random.random()
            tr=tr*r
        x.append(-(math.log10(tr))/a)
    return x
gam=gamma(5,20)

def normal(media,desviacion):
    x=[]
    for _ in range(1000):
        sum = 0.0
        for _ in range (12):
            r=random.random()
            sum += r
        x+=[desviacion*(sum-6.0)+media]
    return x
nor=normal(2.35,30)

def TestChiCuadNormal(nums):
    print("Test de bondad Chi Cuadrado para distribucion Normal")
    observado = []
    esperado = []
    a1=0
    a2=0
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9), 2)
    c =-80
    for _ in range(10):
        x = 0
        for num in nums:
            if  (c-20) <= float(num) <= c : x+=1
        observado.append(x)
        a1+=(c-10)*x
        a2+=((c-10)**2)*x
        c +=20
    a1=a1/1000
    a2=a2/1000
    desviacion=math.sqrt(a2-a1**2)
    media=a1
    c=-80
    esperado=[]
    for i in range(10):
        esperado.append(1000*(norm.cdf((c-media)/desviacion)-norm.cdf(((c-20)-media)/desviacion)))
        c+=20
    x2 = 0
    for i in range(len(observado)):
        x2 += (((observado[i]-esperado[i])**2)/esperado[i])

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

def pascal(k,q):
    nx = []
    for _ in range(1000):
        tr = 1
        qr = math.log10(q)
        for _ in range(k):
            r = random.random()
            tr *= r
        x = int(math.log10(tr)//qr)
        nx.append(x)
    return nx
pas=pascal(5,0.4)

def binomial (n,p):
    x=[]
    for _ in range(1000):
        y=0.0
        for _ in range(1,n):
            r = random.random()
            if (r-p) <0:
                y+=1.0
        x.append(y)
    return x
bino=binomial(1000,0.4)

def TestChiCuadBinomial(nums):
    print("Test de bondad Chi Cuadrado para distribucion Binomial")
    observado = []
    esperado = []
    X=scipy.stats.binom(1000,0.4)
    c =354
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9), 2)
    for _ in range(10):
        x = 0
        for num in nums:
            if  (c-14) <= float(num) <= c : x+=1
        observado.append(x)
        total=sum(X.pmf(k) for k in range (c)) - sum(X.pmf(m) for m in range (c-14))
        esperado.append(1000*total)
        c +=14
    x2 = 0
    for i in range(len(observado)):
        x2 += (((observado[i]-esperado[i])**2)/esperado[i])

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

def hipergeometrica(tn,ns,p):
    x=[]
    for _ in range(1000):
        tn1=tn
        ns1=ns
        p1=p
        y=0.0
        for _ in range(1, ns1):
            r = random.random()
            if(r-p1) > 0:
                s=0.0
            else:
                s=1.0
                y+=1.0
            p1 = (tn1*p1-s)/(tn1-1.0)
            tn1 -= 1.0
        x.append(y)
    return x
hipergeo=hipergeometrica(5000000,500,0.4)

def poisson (lamb):
    x = []
    for _ in range(1000):
        cont=0
        tr=1
        b=0
        while(tr-b >= 0):
            b = math.exp(-lamb)
            r = random.random()
            tr=tr*r
            if(tr-b >= 0):
                cont+=1
        x.append(cont)
    return x
poi=poisson(50)

def TestChiCuadPoisson(nums):
    print("Test de bondad Chi Cuadrado para distribucion Poisson")
    observado = []
    esperado = []
    X=scipy.stats.poisson(50)
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9), 2)
    c =26
    for _ in range(10):
        x = 0
        for j in range(len(nums)):
            if (c-6) <= float(nums[j]) < c:
                x += 1
        observado.append(x)
        total=sum(X.pmf(k) for k in range (c)) -sum(X.pmf(m) for m in range (c-6))
        esperado.append(1000*total)
        c +=6
    x2 = 0
    for i in range(len(observado)):
        x2 += (((observado[i]-esperado[i])**2)/esperado[i])

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

def empirica():
    x=[]
    p=[0.273,0.037,0.195,0.009,0.124,0.058,0.062,0.151,0.047,0.044]
    for _ in range(1000):
        r=random.random()
        a=0
        z=1
        for j in p:
            a+=j
            if (r<=a):
                break
            else:
                z+=1
        x.append(z)
    return x
empi=empirica()

def TestChiCuadEmpirica(nums):
    print("Test de bondad Chi Cuadrado para distribucion Empirica")
    observado = []
    esperado = []
    chiquadesperado = round(scipy.stats.chi2.ppf(1-0.05, 9),2)
    p = [0.273, 0.037, 0.195, 0.009, 0.124, 0.058, 0.062, 0.151, 0.047, 0.044]
    for i in range(10):
        x = 0
        for num in nums:
            if num==i+1:
                x += 1
        observado.append(x)
        esperado.append(1000 * p[i])
    x2 = 0
    for i in range(len(observado)):
        x1 = (((observado[i]-esperado[i])**2)/esperado[i])
        x2 += x1

    print(f"Valor de X2 obtenido = {str(x2)}")
    if (x2 < chiquadesperado): print(f"Valor obtenido: {str(x2)} menor que {str(chiquadesperado)} ,por lo tanto Paso el test")
    else: print("No paso el test")

# GRÁFICOS
def graficar(uni, expo, gam, nor, pas, bino, hipergeo, poi, empi):
    plt.figure(1)
    plt.title("uniforme")
    plt.hist(uni, alpha=.7, edgecolor="black")
    plt.savefig("Uniforme3.png")
    plt.show()

    plt.figure(2)
    plt.title("exponencial")
    plt.hist(expo, 25, alpha=.7, color='g', edgecolor="black")
    plt.savefig("Exponencial3.png")
    plt.show()

    plt.figure(3)
    plt.title("gamma")
    plt.hist(gam,25, alpha=.7, color='r', edgecolor="black")
    plt.savefig("Gamma3.png")
    plt.show()

    plt.figure(4)
    plt.title("normal")
    plt.hist(nor, 25, alpha=.7, color='y', edgecolor="black")
    plt.savefig("Normal3.png")
    plt.show()

    plt.figure(5)
    plt.title("pascal")
    plt.hist(pas, alpha=.7, edgecolor="black", color='cyan')
    plt.savefig("Pascal3.png")
    plt.show()

    plt.figure(6)
    plt.title("binomial")
    plt.hist(bino, 25, alpha=.7, edgecolor="cyan", color='black')
    plt.savefig("Binomial3.png")
    plt.show()

    plt.figure(7)
    plt.title("hipergeometrica")
    plt.hist(hipergeo, 25, alpha=.7, edgecolor="black", color='chocolate')
    plt.savefig("HiperGeometrica3.png")
    plt.show()

    plt.figure(8)
    plt.title("poisson")
    plt.hist(poi, 25, alpha=.7, edgecolor="black", color='orange')
    plt.savefig("Poisson3.png")
    plt.show()

    plt.figure(9)
    plt.title("empirica")
    plt.hist(empi,color='violet', alpha=.7, edgecolor="black",)
    plt.savefig("Empirica3.png")
    plt.show()

#graficar(uni, expo, gam, nor, pas, bino, hipergeo, poi, empi)
# TestChiCuadUni(uni)
# print()
# TestChiCuadExp(expo)
# print()
# TestChiCuadNormal(nor)
# print()
# TestChiCuadBinomial(bino)
# print()
#TestChiCuadPoisson(poi)
# print()
TestChiCuadEmpirica(empi)


# -*- coding: UTF-8 -*-

import os.path
import re
from datetime import datetime,timedelta
from settings import siteinfo,lockfile,usertemppath,isuploadfileoss

def getfiletypename(filetype):
    if (filetype is 1): return '邮箱校验'
    if (filetype is 2): return '短信校验'
    if (filetype is 3): return '密码加密'
    if (filetype is 4): return '金融支付校验'
    if (filetype is 5): return '金融支付强校验'

def file_name(path):
    return os.path.split(path)[1]

def file_extension(filename):
    finfo = os.path.splitext(filename)
    if (len(finfo) < 2):
        return ''
    else:
        return finfo[1]

def lock_site_notify():
    if (os.path.exists(lockfile) is True):
        return True,siteinfo['locknotify']
    return False,siteinfo['sitenotify']

def deep_auth_create(userid, code):
    filename = '%s_%s' % (userid, code)
    nowdir = datetime.now().strftime('%Y%m%d')
    nowpath = os.path.join(usertemppath, nowdir, filename)
    if (os.path.exists(nowpath) is False):
        if (os.path.exists(os.path.join(usertemppath, nowdir)) is False):
            os.makedirs(os.path.join(usertemppath, nowdir))
        with open(nowpath, 'w+') as up:
            up.write(filename)

def deep_auth_check(userid, code):
    filename = '%s_%s' % (userid, code)
    now = datetime.now()
    nowdir = now.strftime('%Y%m%d')
    if (os.path.exists(os.path.join(usertemppath, nowdir, filename)) is True):
        return True
    beforetime = now - timedelta(days=1)
    beforetimedir = beforetime.strftime('%Y%m%d')
    if (os.path.exists(os.path.join(usertemppath, beforetimedir, filename)) is True):
        return True
    return False

def checkfilename(filename):
    filter_rule = '[^\u2E80-\u9FFFa-zA-Z0-9_.]+'
    if (re.search(filter_rule, filename)):
        return False
    else:
        return True

def file_exists(filepath,oss=None):
    if isuploadfileoss is True:
        return oss.exists(filepath)
    else:
        return os.path.exists(filepath)

def file_read(filepath,oss=None):
    if isuploadfileoss is True:
        filedata = oss.Bucket.get_object(filepath).read()
    else:
        with open(filepath, 'r') as f:
            filedata = f.readlines()
    return filedata

def file_write(filepath,filedata,oss=None):
    if isuploadfileoss is True:
        oss.Bucket.put_object(filepath, filedata)
    else:
        with  open(filepath, 'w', encoding='utf-8') as f:
            f.write(filedata)


from collections import deque

# A class to represent a graph object
class Graph:
    # Constructor
    def __init__(self, edges, N):
        # A list of lists to represent an adjacency list
        self.adjList = [[] for _ in range(N)]

        # add edges to the undirected graph
        for (src, dest) in edges:
            self.adjList[src].append(dest)
            self.adjList[dest].append(src)

def DFS(graph, v, discovered):
    # create a stack used to do iterative DFS
    stack = deque()

    # push the source node into the stack
    stack.append(v)

    # loop till stack is empty
    while stack:

        # Pop a vertex from the stack
        v = stack.pop()

        # if the vertex is already discovered yet, ignore it
        if discovered[v]:
            continue

        # we will reach here if the popped vertex `v`
        # is not discovered yet; print it and process
        # its undiscovered adjacent nodes into the stack
        discovered[v] = True
        print(v, end=' ')

        # do for every edge `v —> u`
        adj = graph.adjList[v]
        for i in reversed(range(len(adj))):
            u = adj[i]
            if not discovered[u]:
                stack.append(u)

# Function to perform DFS traversal on the graph on a graph
def recursive_DFS(graph, v, discovered):
    discovered[v] = True            # mark the current node as discovered
    print(v, end=' ')               # print the current node

    # do for every edge `v —> u`
    for u in graph.adjList[v]:
        if not discovered[u]:       # if `u` is not yet discovered
            recursive_DFS(graph, u, discovered)

def discover(graph, numNodes):
    # to keep track of whether a vertex is discovered or not
    discovered = [False] * numNodes

    # Do iterative DFS traversal from all undiscovered nodes to
    # cover all unconnected components of a graph
    for i in range(numNodes):
        if not discovered[i]:
            DFS(graph, i, discovered)
    print()

def discover_recursive(graph, numNodes):
    # to keep track of whether a vertex is discovered or not
    discovered = [False] * numNodes

    # Perform DFS traversal from all undiscovered nodes to
    # cover all unconnected components of a graph
    for i in range(numNodes):
        if not discovered[i]:
            DFS(graph, i, discovered)
    print()

if __name__ == '__main__':

    # List of graph edges as per the above diagram
    edges = [
        # Notice that node 0 is unconnected
        (1, 2), (1, 7), (1, 8), (2, 3), (2, 6), (3, 4),
        (3, 5), (8, 9), (8, 12), (9, 10), (9, 11)
    ]

    # total number of nodes in the graph (0–12)
    N = 13

    # build a graph from the given edges
    graph = Graph(edges, N)

    discover(graph, N)
    discover_recursive(graph, N)


from itertools import tee, chain, islice, zip_longest
import hashlib
from math import log
import re
from contextlib import contextmanager
from .active_timeout import TimeoutExceeded

replacement = re.compile(r'\.0+$')

def intToBytes(num, padding_length=0):
    s = format(num, 'x')
    if len(s) % 2:
        s = '0' + s
    result = bytes.fromhex(s)
    if len(result) < padding_length:
        return (bytes(padding_length-len(result)))+result
    return result

def bytesToInt(num):
    result = 0
    exp = int(256 ** (len(num)-1))
    for i in range(len(num)):
        result += int(num[i]*exp)
        exp //= 256
    return result

def clump(seq, length):
    getter = iter(seq)
    while True:
        try:
            tmp = next(getter)
            yield chain([tmp], islice(getter, length-1))
        except StopIteration:
            return

split_iterable = clump  # for compatability

def _getelem(iterator, elem):
    for row in iterator:
        yield row[elem]

def splice(seq):
    """
    Takes an iterable (which must iterate over a rectangular dataset)
    For an iterator which yields m rows of n elements each
    return n elements of length m, each yielding a different
    column of the original sequence
    """
    getter = iter(seq)
    first = next(getter)
    return [
        _getelem(iterator, i)
        for iterator, i in zip(
            tee(chain([first], getter), len(first)),
            range(len(first))
        )
    ]

def byte_xor(b1, b2):
    return intToBytes(bytesToInt(b1) ^ bytesToInt(b2), max(len(b1), len(b2)))

def hashfile(filepath, algorithm='sha1', length=None):
    reader = open(filepath, mode='rb')
    hasher = hashlib.new(algorithm)
    chunk = reader.read(4096)
    while chunk:
        hasher.update(chunk)
        chunk = reader.read(4096)
    reader.close()
    if isinstance(length, int):
        try:
            return hasher.digest(length)
        except TypeError:
            print(
                "Variable-length digests not supported for",
                algorithm,
                "algorithm on this platform"
            )
    return hasher.digest()

def byteSize(n):
    if n == 0:
        return '0B'
    index = min(8, int(log(abs(n), 1024)))
    suffix = [
        'B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB'
    ][index]
    output = '%.1f' % (n / 1024**index)
    return replacement.sub('', output)+suffix

def first(iterable, predicate=lambda x: True):
    if not callable(predicate):
        def func(item):
            return item == predicate
    else:
        func = predicate
    for item in iterable:
        if func(item):
            return item

class LockTimeoutExceeded(TimeoutExceeded):
    pass

@contextmanager
def context_lock(lock, timeout=-1):
    attempt = lock.acquire(timeout=timeout)
    if not attempt:
        raise LockTimeoutExceeded(
            "Failed to acquire %s after %d seconds" % (lock, timeout)
        )
    try:
        yield lock
    finally:
        lock.release()


import json
import os

for datasets in ["train", "test", "val"]:

    s = ""

    for file_name in os.listdir(r"datasets/jsons"):
        json_path = os.path.join(r"datasets/jsons", file_name)

        if datasets in file_name:

            data = json.load(open(json_path, 'r', encoding='utf-8'), strict=False)
            s = s + "{} {} {} {} {} {} {}\n".format(file_name.replace(".json", ".jpg"), data["start"][0], data["start"][1],
                                                   data["medium"][0], data["medium"][1], data["end"][0], data["end"][1])

    f = open(r"datasets/annotations/{}_dataset.txt".format(datasets), "w", encoding="utf-8")
    f.write(s)
    f.close()


#ref:https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
#nz=100 epoch=100 fid=34 face_reg=80.5%
#nz=100 epoch=100 fid=29 face_reg=86.9% #save Dx higher
#test fid=52.34 face_reg=84.7% original unnormalize
#test fid=51.57 face_reg=88.7% easy unnormalize
#test fid=61.54 face_reg=82.9% transform unnormalize
#test fid=34 face_reg=77.7% easy unnormalize
#test fid=26.87 face_reg=89.8% nz=300 epoch=300 1000/1000
import argparse
import os
import random
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.parallel
from tqdm import tqdm
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from model import Generator,Discriminator
#from IPython.display import HTML
# custom weights initialization called on netG and netD
class UnNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std
    def __call__(self, tensor):
        for t, m, s in zip(tensor, self.mean, self.std):
            t.mul_(s).add_(m)
        return tensor
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
#def main(args):
if __name__ == "__main__":
    # Seed

    # Set random seed for reproducibility
    PATH = "./model/best_model_P1B.pt"
    # ckpt = torch.load(PATH)
    # print(ckpt["epoch"])
    manualSeed = 999
    #manualSeed = random.randint(1, 10000) # use if you want new results
    print("Random Seed: ", manualSeed)
    random.seed(manualSeed)
    torch.manual_seed(manualSeed)

    save_path = "./tmp/"
    # Root directory for dataset
    dataroot = "hw2_data/face"
    # Number of workers for dataloader
    workers = 2
    # Batch size during training
    batch_size = 128
    # Spatial size of training images. All images will be resized to this
    #   size using a transformer.
    image_size = 64
    # Number of channels in the training images. For color images this is 3
    nc = 3
    # Size of z latent vector (i.e. size of generator input)
    nz = 300
    # Size of feature maps in generator
    ngf = 64
    # Size of feature maps in discriminator
    ndf = 64
    # Number of training epochs
    num_epochs = 300
    # Learning rate for optimizers
    lr = 0.0002
    # Beta1 hyperparam for Adam optimizers
    beta1 = 0.5
    # Number of GPUs available. Use 0 for CPU mode.
    ngpu = 1
    # Create the dataset
    mean = (0.5, 0.5, 0.5)
    std = (0.5, 0.5, 0.5)
    dataset = dset.ImageFolder(root=dataroot,transform=transforms.Compose([
                                transforms.Resize(image_size),
                                transforms.CenterCrop(image_size),
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                            ]))
    # Create the dataloader
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                            shuffle=True, num_workers=workers)

    # Decide which device we want to run on
    device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

    # Plot some training images
    real_batch = next(iter(dataloader))
    plt.figure(figsize=(8,8))
    plt.axis("off")
    plt.title("Training Images")
    plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
    plt.savefig("./test.png")

    # Create the generator
    netG = Generator(ngf, nz, nc).to(device)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Apply the weights_init function to randomly initialize all weights
    #  to mean=0, stdev=0.02.
    netG.apply(weights_init)

    # Print the model
    print(netG)

    # Create the Discriminator
    netD = Discriminator(ndf, nc).to(device)
    # Apply the weights_init function to randomly initialize all weights
    #  to mean=0, stdev=0.2.
    netD.apply(weights_init)
    # Print the model
    print(netD)

    # Initialize BCELoss function
    criterion = nn.BCELoss()

    # Establish convention for real and fake labels during training
    real_label = 1.
    fake_label = 0.

    # Setup Adam optimizers for both G and D
    optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
    #optimizerD = optim.SGD(netD.parameters(), lr=lr)
    optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
    transform = transforms.Compose([
            #transforms.Normalize((-1, -1, -1), (2, 2, 2)),
            transforms.ToPILImage(),
        ])
    unorm = UnNormalize(mean, std)
    best_Dx = 0.0
    best_DGz = 0.0
    # Training Loop
    iters = 0
    print("Starting Training Loop...")
    # For each epoch
    for epoch in range(num_epochs):
        # Lists to keep track of progress
        img_list = []
        G_losses = []
        D_losses = []
        # For each batch in the dataloader
        train_loop = tqdm(dataloader,position=0,leave=False,ncols=60,colour='green',desc='Epoch '+str(epoch+1))
        for i, data in enumerate(train_loop, 0):

            ############################
            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
            ###########################
            ## Train with all-real batch
            netD.zero_grad()
            # Format batch
            real_cpu = data[0].to(device)
            b_size = real_cpu.size(0)
            label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
            # Forward pass real batch through D
            output = netD(real_cpu).view(-1)
            # Calculate loss on all-real batch
            errD_real = criterion(output, label)
            # Calculate gradients for D in backward pass
            errD_real.backward()
            D_x = output.mean().item()

            ## Train with all-fake batch
            # Generate batch of latent vectors
            noise = torch.randn(b_size, nz, 1, 1, device=device)
            # Generate fake image batch with G
            fake = netG(noise)
            label.fill_(fake_label)
            # Classify all fake batch with D
            output = netD(fake.detach()).view(-1)
            # Calculate D's loss on the all-fake batch
            errD_fake = criterion(output, label)
            # Calculate the gradients for this batch, accumulated (summed) with previous gradients
            errD_fake.backward()
            D_G_z1 = output.mean().item()
            # Compute error of D as sum over the fake and the real batches
            errD = errD_real + errD_fake
            # Update D
            optimizerD.step()

            ############################
            # (2) Update G network: maximize log(D(G(z)))
            ###########################
            netG.zero_grad()
            label.fill_(real_label)  # fake labels are real for generator cost
            # Since we just updated D, perform another forward pass of all-fake batch through D
            output = netD(fake).view(-1)
            # Calculate G's loss based on this output
            errG = criterion(output, label)
            # Calculate gradients for G
            errG.backward()
            D_G_z2 = output.mean().item()
            # Update G
            optimizerG.step()

            # Save Losses for plotting later
            G_losses.append(errG.item())
            D_losses.append(errD.item())

            iters += 1
        # Output training stats
        print('[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'% (epoch, num_epochs,errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
        if(D_x > best_Dx):
            torch.save({
                'epoch':(epoch+1),
                'netG_state_dict':netG.state_dict(),
                'netD_state_dict':netD.state_dict(),
                'D_x':D_x,
                'D_G_z1':D_G_z1,
                }, PATH)
            best_Dx = D_x
            best_DGz = D_G_z1

    num_generate_ = 5000
    num_result = 20000
    ckpt = torch.load(PATH)
    print(ckpt["epoch"])
    netG = Generator(ngf, nz, nc)
    netD = Discriminator(ndf, nc)
    netG.load_state_dict(ckpt["netG_state_dict"])
    netD.load_state_dict(ckpt["netD_state_dict"])
    netG.to(device)
    netD.to(device)
    with torch.no_grad():
        # Create batch of latent vectors that we will use to visualize
        #  the progression of the generator
        img_prob = []
        for _ in range(num_result//num_generate_):
            fixed_noise = torch.randn(num_generate_, nz, 1, 1, device=device)
            fake = netG(fixed_noise)
            prob = netD(fake)
            fake = fake.detach().cpu()
            prob = prob.detach().cpu()
            #print(fake[0]) (1000,3,64,64)
            #print(prob.size()) (1000,1,1,1)
            #print(prob[0].item()) 0.05
            for j in range(num_generate_):
                img = fake[j] / 2 + 0.5
                im = transform(img)
                img_prob.append((im,prob[j].item()))

        img_prob.sort(key=lambda tup: tup[1],reverse=True)  # sorts in place

        for j in range(1000):
            path = os.path.join(save_path,f"{j}.png")
            im =  img_prob[j][0]
            im.save(path)


''' Functions to make plots.  At present, just plotting spectra.
jrigby, may 2016 '''
from __future__ import print_function

from builtins import range
from jrr import spec
from jrr import mage
from jrr import util
import pandas
from   math import ceil
from   matplotlib import pyplot as plt
from matplotlib.ticker import FuncFormatter
from mpl_toolkits.axes_grid1 import host_subplot
from matplotlib.ticker import MultipleLocator, FormatStrFormatter, MaxNLocator
import numpy as np
from cycler import cycler
from matplotlib.backends.backend_pdf import PdfPages

color1 = 'k'     # color for spectra
color2 = '0.65'  # color for uncertainty spectra
color3 = '0.5'   # color for continuum

def force_axisticks_linear(whichaxis) :
    for axis in whichaxis :  #ax.yaxis]:
        formatter = FuncFormatter(lambda y, _: '{:.16g}'.format(y))
        axis.set_major_formatter(formatter)

def annotate_from_dataframe(df, xcol='x', ycol='y', text='label', xycoords='data', xytext=(4,3), ha='right', va='center', fontsize=10, fontname='Times New Roman', arrowprops=None, ax=None, log=False) :
    # Matplotlib annotate works one annotation at a time -- it can't handle arrays.  Workaround: iterate through the dataframe.
    for index, row in df.iterrows():
        if text=='index' : putlabel = index   # We may want to label the index, and that's handled differently from df['acolumn']
        else             : putlabel = row[text]
        if log == True :  x = np.log10(row[xcol])  ; y = np.log10(row[ycol]) # plot the log of x, y
        else           :  x = row[xcol]            ; y = row[ycol]
        if ax==None:
            plt.annotate(putlabel, xy=(x,y), xycoords=xycoords, xytext=xytext, textcoords="offset points", ha=ha, fontsize=fontsize, \
                             fontname=fontname, arrowprops=arrowprops)
        else:
            ax.annotate(putlabel, xy=(x,y), xycoords=xycoords, xytext=xytext, textcoords="offset points", ha=ha, fontsize=fontsize, \
                            fontname=fontname, arrowprops=arrowprops)
    return(0)

def standard_colors1():
    return(['black', 'blue', 'green', 'purple', 'red', 'orange', 'cyan'])
def standard_colors2():
#    return(['black', 'midnightblue', 'dimgray', 'lightgrey'])
    return(['black', 'blue', 'dimgray', 'lightgrey'])
def standard_colors3():
    return(['blue', 'black', 'red'])
def standard_colors4():
    return([color1, color2, color3])

def boxplot_spectra(wave, fnu, dfnu, line_label, line_center, redshift, win, Ncol, spec_label="",figsize=(8,8), vel_plot=True, verbose=True) :
    '''Make a plot of flux density versus rest-frame velocity or rest-frame wavelength for several spectral lines, in a
    [Nrow x Ncol] box.
    Inputs are:
    wave:        observed wavelength (currently Angstrom)
    fnu:         cont-normalized flux density (erg/s/cm^2/Hz)
    dfnu:        1 sigma uncertainty on fnu
    line_label:  tuple of line labels
    line_center: np array of rest-frame wavelengths (Angstroms)
    redshift:    redshift to assume,
    win:         window to use (km/s if vel_plot; if not, wavelength units.)
    Ncol:        number of columns to plot.  Number of rows set automatically.
    spec_label:  Name of spectrum to plot in lower panel.  May be blank
    figsize:     (Optional) Figure size in inches.
    vel_plot:    (Optional, Boolean):  If True, x-axis is velocity.  If False, x-axis is wavelength.
    '''
    Nrow = ceil(len(line_center) / Ncol)  # Calculate how many rows to generate
    fig = plt.figure(figsize=figsize)
    restwave = wave / (1.0 + redshift)

    for ii, dum in enumerate(line_label) :
        if verbose : print("    Plotting ", line_label[ii], " at ", line_center[ii])
        ax = fig.add_subplot(Nrow, Ncol, ii+1)
        plt.annotate( line_label[ii], (0.3,0.9), xycoords="axes fraction")
        if(vel_plot) :
            vel = spec.convert_restwave_to_velocity(restwave, line_center[ii])   # velocity in km/s
            in_window = vel.between(-1*win, win)
            plt.step(vel[in_window], fnu[in_window], color=color1)   # assumes input is continuum-normalized
            plt.step(vel[in_window], dfnu[in_window], color=color2)  # plot uncertainty
            plt.plot( (-1*win, win), (1.0,1.0), color=color3)        # plot unity continuum
            plt.plot( (0., 0.), (0.0,2), color=color2, linewidth=2)  # plot tics at zero velocity
            #plt.ylim(0.0, 1.5)  # May need to change these limits
            plt.xlim(-1*win, win)
        else :
            in_window = restwave.between((line_center[ii] - win), (line_center[ii] + win))
            plt.step(restwave[in_window], fnu[in_window],  color=color1)
            plt.step(restwave[in_window], dfnu[in_window], color=color2)
            plt.plot( (line_center[ii] - win, line_center[ii] + win), (1., 1.), color=color3)  # plot unity continuum
            plt.plot( (line_center[ii], line_center[ii]), (0.0,2), color=color3, linewidth=2)  # plot tics at zero velocity
            plt.xlim(line_center[ii] - win, line_center[ii] + win)
        if ii == len(line_label) -1 :
            if vel_plot :
                plt.xlabel("rest-frame velocity (km/s)")  # if last subplot, make xlabel
            else :
                plt.xlabel(r'rest-frame wavelength($\AA$)')
        else :
            ax.axes.xaxis.set_ticklabels([])  # if not last subplot, suppress  numbers on x axis
        fig.subplots_adjust(hspace=0)

def boxplot_Nspectra(thewaves, thefnus, thedfnus, thezs, line_label, line_center, spec_label=(), win=2000, Ncol=1, LL=(), extra_label="",figsize=(8,16), vel_plot=True, plot_xaxis=True, ymax=(), colortab=False, verbose=True, drawunity=False, label_loc=(0.55,0.85), lw=(1,),  second_doublet=False, plot_linelabel=True, makenewfig=True) :
    '''Plot flux density versus rest-frame velocity or rest-frame wavelength for several spectral lines,
    in a [Nrow x Ncol] box.  CAN PLOT MULTIPLE SPECTRA IN EACH BOX.
    Inputs are:
    thewaves:        tuple of arrays of observed wavelength (Angstrom).  If only 1 spectrum, use thewaves=(wav_ar,) to keep as tuple.
    thefnus:         tuple of arrays of cont-normalized flux density (erg/s/cm^2/Hz)  ** should this be flambda instead?.
    thedfnus:        tuple of arrays of 1 sigma uncertainty on fnu
    thezs:           tuple of arrays of redshifts, to convert to rest-frame
    line_label:      tuple of line labels.  Makes one box per line
    line_center:     np array of rest-frame wavelengths (Angstroms).  Makes one box per line
    spec_label:      (Optional), labels to name the spectra in a legend
    win:             window to use (km/s if vel_plot; if not, wavelength units.)  Format can be scalar, interpreted as +-win, or list: (-300, 300)
    Ncol:            number of columns to plot.  Number of rows set automatically.
    LL:              (Optional) Linelist, a Pandas data frame, read in by jrr.mage.get_linelist(linelist), to mark other lines
    theconts:        (Optional, tuple of continuum fits)
    extra_label:     (Optional) Extra text to annotate to lower panel.
    figsize:         (Optional) Figure size in inches.
    vel_plot:        (Optional, Bool) If True, x-axis is velocity.  If False, x-axis is wavelength.
    plot_xaxis:      (Optional, Bool) Plot the xaxis?
    ymax:            (Optional), array of max y values to use in subplots.  Over-rides auto-scaling. Size must == line_center
    colortab:        (Optional) color table to use, to replace default colortab
    verbose:         (Optional) Verbose or terse account of status?
    drawunity:       (Optional), draw a line at unity?  Good for spectra that are already continuum-normalized
    label_loc:       (Optional), where to draw the line label for each boxplot.  By default, upper right
    lw:              (Optional), scalar or tuple of linewidths.  If scalar, same lw for all
    second_doublet:  (Optional)  a second wavelength to draw a vertical line over.  Useful for the second transition of a doublet.
    plot_linelabel:  (Optional) flag to suppress plotting of line_label
    thewaves is now a tuple? of wavelength arrays.  same for thefnus, the dfnus, thezs
    If only plotting one spectrum, still need input format to be tuples, as in thewaves=(wave_array,).'''

    if colortab :  mycol=colortab
    else : mycol = standard_colors1()
    linestyles = ['solid'] #, 'dashed', 'dotted']#, 'dashdot']
    plt.rc('axes', prop_cycle=(cycler('color', mycol) * cycler('linestyle', linestyles)))
    if type(line_center) == np.float64 :  Nrow=1
    else :  Nrow = int(np.ceil( float(len(line_center)) / Ncol))  # Calculate how many rows to generate
    #print "DEBUGGING ", Nrow, Ncol, len(line_center)
    if makenewfig: fig = plt.figure(figsize=figsize) # turn this off if calling this to make a grid

    if len(spec_label) == len(thewaves) :  pass  #print "DEBUGGING, adding spectrum label to legend"
    else : spec_label = np.repeat('_nolegend_', len(thewaves))  # override the legend

    if hasattr(win, "__len__") and len(win) == 2:
        win0 = win[0]  ; win1 = win[1]
    elif not hasattr(win, "__len__")  :    win0 = -1*win ; win1=win
    else : raise Exception("ERROR: Window width win must either be a scalar, or a two- element list/tuple/array")

    if len(lw) == 1 : lw = lw * np.ones(shape=len(thewaves))  # if lw is a scalar, change to an array, same for all
    elif len(lw) != len(thewaves) : raise Exception('Error, linewidth lw must be scalar (len=1) or tuple with same sizes as thewaves')
    for ii, dum in enumerate(line_label) :
        if verbose : print("    Plotting ", line_label[ii], " at ", line_center[ii])
        ax = fig.add_subplot(Nrow, Ncol, ii+1)
        if plot_linelabel: plt.annotate( line_label[ii], label_loc, xycoords="axes fraction")
        max_in_window = 0.
        if(vel_plot) :
            for ss in range(0, len(thewaves)) :  # For each spectrum to overplot
                restwave = thewaves[ss] / (1.0 + thezs[ss])
                vel = spec.convert_restwave_to_velocity(restwave, line_center[ii])   # velocity in km/s
                in_window = vel.between(win0, win1)
                plt.step(vel[in_window], thefnus[ss][in_window], color=mycol[ss], linewidth=lw[ss], label=spec_label[ss])
                if len(thedfnus[ss]) :
                    #plt.step(vel[in_window], thedfnus[ss][in_window], color=mycol[ss], linewidth=1, label='_nolegend_')  # plot uncertainty
                    plt.step(vel[in_window], thedfnus[ss][in_window], color='grey', linewidth=1, label='_nolegend_')  # plot uncertainty in grey
                thismax = thefnus[ss][in_window].max()
                max_in_window =  util.robust_max((thismax, max_in_window))
            plt.plot( (0., 0.), (0.0,100), color=color2, linewidth=3)  # plot tics at zero velocity
            if drawunity: plt.plot( (win0, win1), (1.0,1.0), color=color3)        # plot unity continuum
            plt.xlim(win0, win1)
        else :
            for ss in range(0, len(thewaves)) :  # For each spectrum to overplot
                restwave = thewaves[ss] / (1.0 + thezs[ss])
                in_window = restwave.between((line_center[ii] + win0), (line_center[ii] + win1))
                plt.step(restwave[in_window], thefnus[ss][in_window], color=mycol[ss], linewidth=lw[ss], label=spec_label[ss])
                if len(thedfnus[ss]) :
                    #plt.step(restwave[in_window], thedfnus[ss][in_window], color=mycol[ss], linewidth=1, label='_nolegend_')
                    plt.step(restwave[in_window], thedfnus[ss][in_window], color='grey', linewidth=1, label='_nolegend_')
                thismax = thefnus[ss][in_window].max()
                max_in_window =  util.robust_max((thismax, max_in_window))
            plt.plot( (line_center[ii], line_center[ii]), (0.0,100), color=mycol[1], linewidth=3, alpha=0.4)  # plot tics at zero velocity
            if second_doublet:  plt.plot( (second_doublet, second_doublet), (0.0,100), color=mycol[4], linewidth=3, alpha=0.4) # plot the second component of the doublet
            if drawunity: plt.plot( (line_center[ii] + win0, line_center[ii] + win1), (1.0,1.0), color=color3,  label='_nolegend_')  # plot unity continuum
            plt.xlim(line_center[ii] +win0, line_center[ii] + win1)
        ax.locator_params(axis='y', nbins=4)
        ax.locator_params(axis='x', nbins=7)
        if len(ymax)  == 1 :  plt.ylim(0, ymax[0])   # user can over-ride autoscaling
        elif len(ymax) > 1 :  plt.ylim(0, ymax[ii])  # user can over-ride autoscaling
        else :                plt.ylim(0., max_in_window*1.1)
        if len(LL) :
            mage.plot_linelist(LL, thezs[0], True, vel_plot, line_center[ii])  # plot the line IDs for the first spectrum only
        if ii == len(line_label) -1 :
            plt.annotate(extra_label, (0.6,0.1), xycoords="axes fraction")
            if vel_plot :
                plt.xlabel("rest-frame velocity (km/s)")  # if last subplot, make xlabel
            else :
                plt.xlabel(r'rest-frame wavelength ($\rm \AA$)')
        if (not plot_xaxis) and (ii < len(line_label)-1) :
            ax.axes.xaxis.set_ticklabels([])  # if not last subplot, suppress  numbers on x axis
            # But warning, this will disable x= on interactive matplotlib.  comment out above line to measure numbers interactively on graph
        if not plot_xaxis :
            fig.subplots_adjust(hspace=0)

def velocity_overplot(wave, fnu, line_label, line_center, redshift, vwin1, vwin2, figsize=(8,8), colortab=False) :
    '''Same as boxplot_spectra(vel_plot=True), but OVERPLOT all the lines on 1 plot, no subplots'''
    if colortab :  mycol=colortab
    else : mycol = standard_colors1()
    restwave = wave / (1.0 + redshift)
    plt.figure(figsize=figsize)
    for ii, dum in enumerate(line_label) :
        print(ii, "Trying to plot ", line_label[ii], " at ", line_center[ii])
        vel = spec.convert_restwave_to_velocity(restwave, line_center[ii])   # velocity in km/s
        in_window = vel.between((vwin1), vwin2)
        plt.step(vel[in_window], fnu[in_window], color=mycol[ii], label=line_label[ii])
        plt.xlim(vwin1, vwin2)
        plt.ylim(0, 1.4)
        xy = (0.1,  0.26 - float(ii)/len(line_label)*0.3)
        #print "DEBUG, xy is ", xy
        #plt.annotate( line_label[ii], xy, xycoords="axes fraction", color=hist.get_color())
        plt.legend()
        plt.xlabel("rest-frame velocity (km/s)")
    plt.plot( (0.,0.), plt.ylim(),  color=color2, linewidth=2)  # plot tics at zero velocity
    plt.plot( plt.xlim(), (1.0,1.0), color=color2)  # plot unity continuum

def annotate_echelle() :
    plt.annotate("f_nu (erg/s/cm^2/Hz)",  xy=(0.05,0.055), color="black", xycoords="figure fraction")
    plt.annotate("Line color coding:", xy=(0.05,0.043), color="black", xycoords="figure fraction")
    plt.annotate("Photosphere", xy=(0.05,0.03), color="blue", xycoords="figure fraction")
    plt.annotate("Emission",    xy=(0.17,0.03), color="red", xycoords="figure fraction")
    plt.annotate("ISM",         xy=(0.25,0.03), color="green", xycoords="figure fraction")
    plt.annotate("Wind",        xy=(0.32,0.03), color="cyan", xycoords="figure fraction")
    plt.annotate("Fine structure", xy=(0.4,0.03), color="purple", xycoords="figure fraction")
    plt.annotate("Intervening", xy=(0.55,0.03), color="orange", xycoords="figure fraction")

def echelle_spectrum(the_dfs, the_zzs, LL=(), Npages=4, Npanels=24, plotsize=(11,16), outfile="multipanel.pdf", title="", norm_by_cont=False, plot_cont=False, apply_bad=False, waverange=(), colwave='wave', colfnu='fnu', colfnu_u='fnu_u', colcont='fnu_autocont', colortab=False, plotx2=True, ylim=(), topfid=(1.1,0.3), annotate=annotate_echelle, verbose=False) :
    ''' Plot an echelle(ette) spectrum with several pages and many panels for page.  Based on multipanel-spectrum-ids.py,
    but modular. and with more features.  Inputs:
    the_dfs:    an array of dataframes containing spectra, to plot. Usually this is just one dataframe, one spectrum,
                but want to allow overplotting of multiple spectra.  First spectrum will be used to set the wavelength range.
    the_zzs:    an array of redshifts
    LL:         (Optional). a linelist to plot features.  At present, takes redshift from the_zzs[0].
    Npages:     pages of plots in the output PDF.  Will split evenly.
    Npanels:    Total number of panels to plot, across all pages
    plotsize:   (Optional) Size of the plot
    outfile:    Name of output PDF file.
    norm_by_cont: (Optional), Normalize by continuum?
    plot_cont:  (Optional) Plot the continuum, too?
    apply_bad:  (Optional) Apply bad pixel map?
    waverange:  (Optional) Only plot this wavelength range, rather than full range. Format is (wave_lo, wave_hi).
    colwave, colfnu, colfnu_u, colcont:  columns in the dataframes to use for wave, fnu, fnu_uncert, continuum
    colortab:   (Optional) color table to use, to replace default colortab
    plotx2:     Plot upper x-axis, the rest-frame wavelength?  Turn off when lower x is already rest-frame
    ylim:       (Optional), over-ride my auto-sensing of the ylims for each subplot
    topfid:     (Optional) fiddle parameters to adjust tightness of y axis.  first scales the median, second scales the IQR
    annotate:   (Optional) Function that will annotate the first page.  For plot customization.'''

    contcolor = '0.75'
    if colortab :  linecolor=colortab
    else : linecolor = standard_colors2()
    pp = PdfPages(outfile)  # output
    # Set up the wavelength ranges.  Each panel has same # of pixels, not same dlambda, since dispersion may change.
    if not len(waverange) in (0, 2) : raise Exception("Error: len(waverange) must be 0 or 2")
    if len(waverange) :  # If user selected wavelength range, modify index, but keep same # of pixels per panel
        first_index = the_dfs[0][the_dfs[0][colwave].between(waverange[0], waverange[0]*1.01)][colwave].index[0]
        last_index  = the_dfs[0][the_dfs[0][colwave].between(waverange[1], waverange[1]*1.01)][colwave].index[-1]
    else :
        first_index = 0
        last_index  = len(the_dfs[0][colwave])-1
    ind = (np.linspace(first_index, last_index, num=Npanels+1)).astype(np.int)
    start = the_dfs[0][colwave][ind[0:-1]].values
    end   = the_dfs[0][colwave][ind[1:]].values

    max_per_page = np.int(np.float64(Npanels) / Npages)
    the_format = max_per_page*100 + 11  # makes label: 311 for 3 rows, 1 col, start at 1

    for df in the_dfs :
        if norm_by_cont :   df['normby'] = df[colcont]  # easy to plot whether or not it's continuum normalized
        else :              df['normby'] = 1.0
    for kk in range(0, Npanels) :
        if kk % max_per_page == 0 :   # Start a new page
            current_page = kk /  max_per_page + 1  # should be int division
            print("Starting page", current_page, "of", Npages)
            fig    = plt.figure(figsize=plotsize)
        subit = host_subplot(the_format + kk % max_per_page)
        if verbose : print("Plotting panel", kk, "for wavelengths", start[kk], "to",end[kk])
        for ss, df in enumerate(the_dfs) :
            if apply_bad : subset = df[df[colwave].between(start[kk], end[kk]) & ~df['badmask']]
            else:          subset = df[df[colwave].between(start[kk], end[kk])]

            plt.plot(subset[colwave], subset[colfnu]/subset['normby'],   color=linecolor[ss],  linestyle='steps', linewidth=1)
            plt.plot(subset[colwave], subset[colfnu_u]/subset['normby'], color=linecolor[ss], linestyle='steps', linewidth=1) # plot 1 sigma uncert spectrum
            if plot_cont :  plt.plot(subset[colwave], subset[colcont].astype(np.float)/subset['normby'], contcolor, linestyle='steps', zorder=1, linewidth=1) # plot the continuum
            if ss == 0 :  # If the first df, set the plot ranges
                top = (subset[colfnu]/subset['normby']).median()*topfid[0] + util.IQR(subset[colfnu]/subset['normby'])*topfid[1]
                #print "DEBUGGING top", kk, top, (subset[colfnu]/subset['normby']).median(), + util.IQR(subset[colfnu]/subset['normby'])
        if len(ylim) == 2:   plt.ylim(ylim[0], ylim[1])
        else:                plt.ylim(0, top)  # trying to fix weird autoscaling from bad pixels
        plt.xlim(start[kk], end[kk])
        if len(LL) : mage.plot_linelist(LL, the_zzs[0])   # Plot the line IDs.  Use the first redshift for z.
        if plotx2:
            upper = subit.twiny()  # make upper x-axis in rest wave (systemic, or of abs lines)
            upper.set_xlim( start[kk]/(1.0+the_zzs[0]), end[kk]/(1.0+the_zzs[0]))
            upper.locator_params(axis='x', nbins=5)
            subit.locator_params(axis='x', nbins=5)
            subit.xaxis.tick_bottom()  # don't let lower ticks be mirrored  on upper axis

        if  kk % max_per_page == (max_per_page-1) or kk == Npanels-1:   # last plot on this page
            if plotx2 : subit.set_xlabel(r"observed-frame vacuum wavelength ($\AA$)")  ## COMMENDED OUT FOR mage stack paper
            else:       subit.set_xlabel(r"rest-frame vacuum wavelength ($\AA$)")  # Just for mage stack paper
            #plt.ylabel('fnu') # fnu in cgs units: erg/s/cm^2/Hz
            plt.ylabel('relative flux') # **temp for stacked paper
            pp.savefig(bbox_inches='tight')
            #fig.canvas.draw()
        if  kk % max_per_page == 0 and plotx2 :  # first plot on the page
            upper.set_xlabel(u"rest-frame vacuum wavelength ($\AA$)")

        if(kk == 0):  # first page
            if callable(annotate) : annotate()
            plt.suptitle(title)  # global title
    pp.close()
    print("   Generated PDF:  ", outfile)
    return(0)


#!/usr/bin/python3
import pygal
import sys, json;
from datetime import datetime
import platform

step = 2
data = json.load(sys.stdin)
classes = set([n["device_class"] for n in data["nodes"]])

hist = pygal.Histogram(title="%s on %s" % (str(datetime.now()), platform.node()), x_title="Utilization %", y_title="# of OSDs")
for c in sorted(classes):
	bins = [[x, 0] for x in range(0,100,step)]
	for node in data["nodes"]:
		if node["device_class"] == c:
			u = (int)(node["utilization"]/step);
			bins[u][1] += 1;
	hist.add(c, [(b[1], b[0], b[0]+step) for b in bins])

hist.render_to_file('/dev/stdout')


#######################################################
# Script with unit tests for CYMDISTToFMU
#
# TSNouidui@lbl.gov                          2016-09-06
#######################################################
import unittest
import os
import sys
import platform
import subprocess
import shutil
from datetime import datetime

# Appending parser_path to the system path os required to be able
# to find the CYMDISTToFMU model from this directory
script_path = os.path.dirname(os.path.realpath(__file__))
parser_path = os.path.abspath(os.path.join(script_path, '..', 'parser'))
sys.path.append(parser_path)

import CYMDISTToFMU as cymdist

XSD_SCHEMA = 'CYMDISTModelDescription.xsd'
NEEDSEXECUTIONTOOL = 'needsExecutionTool'
MODELDESCRIPTION = 'modelDescription.xml'
CYMDISTModelicaTemplate_MO = 'CYMDISTModelicaTemplate.mo'
CYMDISTModelicaTemplate_Dymola_MOS = 'CYMDISTModelicaTemplate_Dymola.mos'
CYMDISTModelicaTemplate_JModelica_MOS = 'CYMDISTModelicaTemplate_JModelica.py'
CYMDISTModelicaTemplate_OpenModelica_MOS = 'CYMDISTModelicaTemplate_OpenModelica.mos'
XML_MODELDESCRIPTION = 'CYMDISTModelDescription.xml'
# Get the path to the templates files
script_path = os.path.dirname(os.path.realpath(__file__))
utilities_path = os.path.join(script_path, '..', 'parser', 'utilities')
PYTHON_SCRIPT_PATH = os.path.join(utilities_path, 'cymdist_wrapper.py')
MO_TEMPLATE_PATH = os.path.join(utilities_path, CYMDISTModelicaTemplate_MO)
MOS_TEMPLATE_PATH_DYMOLA = os.path.join(
    utilities_path, CYMDISTModelicaTemplate_Dymola_MOS)
MOS_TEMPLATE_PATH_JMODELICA = os.path.join(
    utilities_path, CYMDISTModelicaTemplate_JModelica_MOS)
MOS_TEMPLATE_PATH_OPENMODELICA = os.path.join(
    utilities_path, CYMDISTModelicaTemplate_OpenModelica_MOS)
XSD_FILE_PATH = os.path.join(utilities_path, XSD_SCHEMA)
XML_INPUT_FILE = os.path.join(utilities_path, XML_MODELDESCRIPTION)
CYMDISTToFMU_LIB_PATH = os.path.join(
    script_path, '..', 'parser', 'libraries', 'modelica')
python_scripts_path = [PYTHON_SCRIPT_PATH]

if(platform.system().lower() == 'windows'):
    python_scripts_path = [item.replace('\\', '\\\\') for item in [
        PYTHON_SCRIPT_PATH]]

CYMDIST_T = cymdist.CYMDISTToFMU('con_path',
                                       XML_INPUT_FILE,
                                       CYMDISTToFMU_LIB_PATH,
                                       MO_TEMPLATE_PATH,
                                       MOS_TEMPLATE_PATH_DYMOLA,
                                       XSD_FILE_PATH,
                                       '34',
                                       python_scripts_path,
                                       '2.0',
                                       'me',
                                       'dymola',
                                       None,
                                       None,
                                       'false')

class Tester(unittest.TestCase):
    '''
    Class that runs all regression tests.

    '''

    def find_executable(self, tool):

        '''
        Function for checking if Dymola, JModelica, or OpenModelica is installed.

        '''

        if tool == 'jmodelica' and platform.system().lower() == "windows":
            tool = 'pylab'
        if tool == 'jmodelica' and platform.system().lower() == "linux":
            tool = 'jm_python.sh'

        if tool == 'openmodelica' and platform.system().lower() == "windows":
            tool = 'omc'

        cmd = "where" if platform.system() == "Windows" else "which"
        try:
            return subprocess.call([cmd, tool])
        except:
            print ("No executable for tool={!s}".format(tool))
            return 1

    def run_cymdist (self, tool):

        '''
        Function for running FMUs exported from Dymola, JModelica, and OpenModelica with PyFMI.

        '''

        try:
            from pyfmi import load_fmu
        except BaseException:
            print ('PyFMI not installed. Test will not be be run.')
            return
        if (tool=='openmodelica' and platform.system().lower() == 'linux'):
                print ('tool={!s} is not supported on Linux'.format(tool))
                return

        else:
        # Export FMUs which are needed to run the cases.
            if tool == 'openmodelica':
                modPat = 'OPENMODELICALIBRARY'
                mosT = MOS_TEMPLATE_PATH_OPENMODELICA
            elif tool == 'dymola':
                modPat = 'MODELICAPATH'
                mosT = MOS_TEMPLATE_PATH_DYMOLA
            elif tool == 'jmodelica':
                # Unset environment variable
                if (os.environ.get('MODELICAPATH') is not None):
                    del os.environ['MODELICAPATH']
                modPat = None
                mosT = MOS_TEMPLATE_PATH_JMODELICA
            for version in ['2']:
                if (tool == 'openmodelica' or tool == 'jmodelica'):
                    version = str(float(version))
                for api in ['me']:
                    if (tool == 'openmodelica' and version == '1.0' and api == 'cs'):
                        print (
                            'tool={!s} with FMI version={!s} and FMI API={!s} is not supported.'.format(
                                tool, version, api))
                        continue
                    for cs_xml in ['true']:
                        if (version == '1'):
                            continue
                        CYMDIST_Test = cymdist.CYMDISTToFMU(
                            'con_path',
                            XML_INPUT_FILE,
                            CYMDISTToFMU_LIB_PATH,
                            MO_TEMPLATE_PATH,
                            mosT,
                            XSD_FILE_PATH,
                            '34',
                            python_scripts_path,
                            version,
                            api,
                            tool,
                            None,
                            modPat,
                            cs_xml)

                        print (
                            'Export the cymdist with tool={!s}, FMI version={!s}, FMI API={!s}'.format(
                                tool, version, api))
                        start = datetime.now()
                        # CYMDIST_Test.print_mo()
                        # CYMDIST_Test.generate_fmu()
                        # CYMDIST_Test.clean_temporary()
                        # CYMDIST_Test.rewrite_fmu()
                        end = datetime.now()
                        print(
                            'Export CYMDIST as an FMU in {!s} seconds.'.format(
                                (end - start).total_seconds()))

                        fmu_path = os.path.join(
                        script_path, '..', 'fmus', tool, platform.system().lower())
                        print(
                            'Copy CYMDIST.fmu to {!s}.'.format(fmu_path))
                        shutil.copy2('CYMDIST.fmu', fmu_path)

        fmu_path = os.path.join(
                script_path, '..', 'fmus', tool, platform.system().lower(), 'CYMDIST.fmu')
        # Parameters which will be arguments of the function
        start_time = 0.0
        stop_time = 0.1

        print ('Starting the simulation with {!s}'.format(tool))
        start = datetime.now()

        cymdist_input_valref = []
        cymdist_output_valref = []

        sim_mod = load_fmu(fmu_path, log_level=7)
        sim_mod.setup_experiment(
            start_time=start_time, stop_time=stop_time)

        cymdist_con_val_str=os.path.abspath('Z:\\thierry\\proj\\cyder_repo\\NO_SHARING\\CYMDIST\\config.json')
        if sys.version_info.major>2:
            cymdist_con_val_str=bytes(cymdist_con_val_str,'utf-8')
        cymdist_con_val_ref=sim_mod.get_variable_valueref('_configurationFileName')

        # Define the inputs
        cymdist_input_names = ['VMAG_A', 'VMAG_B', 'VMAG_C', 'VANG_A', 'VANG_B', 'VANG_C']
        cymdist_input_values = [2520, 2520, 2520, 0, -120, 120]
        cymdist_output_names = ['IA', 'IAngleA', 'IB', 'IAngleB', 'IC', 'IAngleC']

        # Get the value references of cymdist inputs
        for elem in cymdist_input_names:
            cymdist_input_valref.append(
                sim_mod.get_variable_valueref(elem))

        # Get the value references of cymdist outputs
        for elem in cymdist_output_names:
            cymdist_output_valref.append(
                sim_mod.get_variable_valueref(elem))

        # Set the flag to save the results
        sim_mod.set('_saveToFile', 0)

        if not (tool=="jmodelica"):
            sim_mod.set_string([cymdist_con_val_ref], [cymdist_con_val_str])

        # Initialize the FMUs
        sim_mod.initialize()

        # Call event update prior to entering continuous mode.
        sim_mod.event_update()

        # Enter continuous time mode
        sim_mod.enter_continuous_time_mode()

        sim_mod.set_real(cymdist_input_valref, cymdist_input_values)

        end = datetime.now()

        print(
            'Ran a single CYMDIST simulation with {!s} FMU={!s} in {!s} seconds.'.format(
                tool, fmu_path, (end - start).total_seconds()))
        if not (tool=='openmodelica'):
            #print(sim_mod.get_real(cymdist_output_valref))
            # PyFMI fails to get the output of an OpenModelica FMU
            if(abs(sum(sim_mod.get_real(cymdist_output_valref))-618.57)>1e-3):
                raise ValueError('Values are not matching.')
        # Terminate FMUs
        sim_mod.terminate()

    def test_check_duplicates(self):
        '''
        Test the function check_duplicates().

        '''

        # Array does not contain duplicates variables.
        cymdist.check_duplicates(['x1', 'x2', 'x3', 'x4'])

        # Array contain duplicates variables.
        with self.assertRaises(AssertionError):
            cymdist.check_duplicates(['x1', 'x1', 'x3', 'x4'])

    def test_sanitize_name(self):
        '''
        Test the function sanitize_name().

        '''

        # Testing name conversions.
        name = cymdist.sanitize_name('test+name')
        self.assertEqual(name, 'test_name', 'Names are not matching.')

        name = cymdist.sanitize_name('0test+*.name')
        self.assertEqual(name, 'f_0test___name', 'Names are not matching.')

    def test_xml_validator(self):
        '''
        Test the function xml_validator().

        '''

        # Testing validation of xml file
        CYMDIST_T.xml_validator()

    def test_xml_parser(self):
        '''
        Test the function xml_validator().

        '''

        # Testing validation of xml file
        CYMDIST_T.xml_parser()

    def test_print_mo(self):
        '''
        Test the function print_mo().

        '''

        # Testing function to print Modelica model.
        CYMDIST_T.print_mo()

    def test_cymdist_to_fmu(self):
        '''
        Test the export of an FMU with various options.

        '''

        for tool in  ['dymola', 'jmodelica', 'openmodelica']:
            retVal=self.find_executable(tool)
            if ((retVal is not None) and retVal!=1):
                print("======tool={!s} was found. Unit Test will be run".format(tool))
            else:
                continue

            if (platform.system().lower() == 'linux' and tool == 'openmodelica'):
                print ('tool={!s} is not supported on Linux.'.format(tool))
                continue
            if tool == 'openmodelica':
                modPat = 'OPENMODELICALIBRARY'
                mosT = MOS_TEMPLATE_PATH_OPENMODELICA
            elif tool == 'dymola':
                modPat = 'MODELICAPATH'
                mosT = MOS_TEMPLATE_PATH_DYMOLA
            elif tool == 'jmodelica':
                if (os.environ.get('MODELICAPATH') is not None):
                    del os.environ['MODELICAPATH']
                modPat = None
                mosT = MOS_TEMPLATE_PATH_JMODELICA
            for version in ['1', '2']:
                if (tool == 'openmodelica' or tool == 'jmodelica'):
                    version = str(float(version))
                for api in ['me', 'cs']:
                    if (tool == 'openmodelica' and version == '1.0' and api == 'cs'):
                        print (
                            'tool={!s} with FMI version={!s} and FMI API={!s} is not supported.'.format(
                                tool, version, api))
                        continue
                    for cs_xml in ['true']:
                        if (version == '1'):
                            continue
                        CYMDIST_Test = cymdist.CYMDISTToFMU(
                            'con_path',
                            XML_INPUT_FILE,
                            CYMDISTToFMU_LIB_PATH,
                            MO_TEMPLATE_PATH,
                            mosT,
                            XSD_FILE_PATH,
                            '34',
                            python_scripts_path,
                            version,
                            api,
                            tool,
                            None,
                            modPat,
                            cs_xml)

                        print (
                            'Export the cymdist with tool={!s}, FMI version={!s}, FMI API={!s}'.format(
                                tool, version, api))
                        start = datetime.now()
                        CYMDIST_Test.print_mo()
                        CYMDIST_Test.generate_fmu()
                        CYMDIST_Test.clean_temporary()
                        CYMDIST_Test.rewrite_fmu()
                        end = datetime.now()
                        print(
                            'Export CYMDIST as an FMU in {!s} seconds.'.format(
                                (end - start).total_seconds()))

    def test_updates_fmu(self):
        '''
        Test the export and updates of FMUs.

        '''

        for tool in  ['dymola', 'jmodelica', 'openmodelica']:
            retVal=self.find_executable(tool)
            if ((retVal is not None) and retVal!=1):
                print("======tool={!s} was found. Unit Test will be run".format(tool))
            else:
                continue

            if (platform.system().lower() == 'linux' and tool == 'openmodelica'):
                print ('tool={!s} is not supported on Linux.'.format(tool))
                continue
            if tool == 'openmodelica':
                modPat = 'OPENMODELICALIBRARY'
                mosT = MOS_TEMPLATE_PATH_OPENMODELICA
            elif tool == 'dymola':
                modPat = 'MODELICAPATH'
                mosT = MOS_TEMPLATE_PATH_DYMOLA
            elif tool == 'jmodelica':
                if (os.environ.get('MODELICAPATH') is not None):
                    del os.environ['MODELICAPATH']
                modPat = None
                mosT = MOS_TEMPLATE_PATH_JMODELICA
            for version in ['2']:
                if (tool == 'openmodelica' or tool == 'jmodelica'):
                    version = str(float(version))
                for api in ['me']:
                    if (tool == 'openmodelica' and version == '1.0' and api == 'cs'):
                        print (
                            'tool={!s} with FMI version={!s} and FMI API={!s} is not supported.'.format(
                                tool, version, api))
                        continue
                    for cs_xml in ['true']:
                        if (version == '1'):
                            continue
                        CYMDIST_Test = cymdist.CYMDISTToFMU(
                            'con_path',
                            XML_INPUT_FILE,
                            CYMDISTToFMU_LIB_PATH,
                            MO_TEMPLATE_PATH,
                            mosT,
                            XSD_FILE_PATH,
                            '34',
                            python_scripts_path,
                            version,
                            api,
                            tool,
                            None,
                            modPat,
                            cs_xml)

                        print (
                            'Export the cymdist with tool={!s}, FMI version={!s}, FMI API={!s}'.format(
                                tool, version, api))
                        start = datetime.now()
                        CYMDIST_Test.print_mo()
                        CYMDIST_Test.generate_fmu()
                        CYMDIST_Test.clean_temporary()
                        CYMDIST_Test.rewrite_fmu()
                        end = datetime.now()
                        print(
                            'Export CYMDIST as an FMU in {!s} seconds.'.format(
                                (end - start).total_seconds()))
                        fmu_path = os.path.join(
                        script_path, '..', 'fmus', tool, platform.system().lower())
                        print(
                            'Copy CYMDIST.fmu to {!s}.'.format(fmu_path))
                        shutil.copy2('CYMDIST.fmu', fmu_path)

    def test_run_cymdist_all(self):
        '''
        Test the execution of one CYMDIST FMU.

        '''

        # Export FMUs which are needed to run the cases.
        for tool in  ['dymola', 'jmodelica', 'openmodelica']:
            retVal=self.find_executable(tool)
            if ((retVal is not None) and retVal!=1):
                print("======tool={!s} was found. Unit Test will be run.".format(tool))
                print('=======The unit test will be run for tool={!s}.'.format(tool))
                self.run_cymdist (tool)
            else:
                continue

    def test_run_cymdist_dymola(self):
        '''
        Test the execution of one CYMDIST FMU.

        '''

        retVal=self.find_executable('dymola')
        if ((retVal is not None) and retVal!=1):
            print("======tool=dymola was found. Unit Test will be run.")
            print('=======The unit test will be run for Dymola.')
            print('=======Make sure that Dymola is on the System Path otherwise the simulation will fail.')
            # Run the cases
            self.run_cymdist ('dymola')
        else:
            return

    def test_run_cymdist_jmodelica(self):
        '''
        Test the execution of one CYMDIST FMU.

        '''

        retVal=self.find_executable('jmodelica')
        if ((retVal is not None) and retVal!=1):
            print("======tool=jmodelica was found. Unit Test will be run.")
            print('=======The unit test will be run for JModelica.')
            print('=======Make sure that JModelica is on the System Path otherwise the simulation will fail.')
            # Run the cases
            self.run_cymdist ('jmodelica')
        else:
            return

    def test_run_cymdist_openmodelica(self):
        '''
        Test the execution of one CYMDIST FMU.

        '''

        retVal=self.find_executable('openmodelica')
        if ((retVal is not None) and retVal!=1):
            print("======tool=openmodelica was found. Unit Test will be run.")
            print('=======The unit test will be run for OpenModelica.')
            print('=======Make sure that OpenModelica is on the System Path otherwise the simulation will fail.')
            # Run the cases
            self.run_cymdist ('openmodelica')
        else:
            return

if __name__ == "__main__":
        # Check command line options
    if (platform.system().lower() in ['windows']):
        unittest.main()
    else:
        print('=========CYMDISTToFMU is only supported Windows.')


# Python imports
from dataclasses import dataclass
from pathlib import Path
import pandas as pd
import numpy as np
import spatialmath
from kincalib.utils.ExperimentUtils import separate_markerandfiducial
from kincalib.utils.Frame import Frame

np.set_printoptions(precision=4, suppress=True, sign=" ")

def get_wrist_fiducials_cp(robot_cp):
    fiducials_cp = robot_cp.loc[robot_cp["m_t"] == "f"]
    steps = fiducials_cp["step"].unique().reshape((-1, 1))
    _, wrist_fiducials = separate_markerandfiducial(
        None, Path("./share/custom_marker_id_112.json"), df=fiducials_cp
    )
    wrist_fiducial_cp = pd.DataFrame(
        np.hstack((steps, wrist_fiducials.T)), columns=["step", "x", "y", "z"]
    )
    return wrist_fiducial_cp

@dataclass
class FRE:
    """Calculate the Fiducial registration error (FRE) between two point clouds. The input data can be specified
    in joints space or cartesian space. If joints data is given, the class requires also a forward kinematic function.
    Each set of points should be passed as a dictionary with the following keys:

    ```dict(mode='',data=df, fk=fk)```

    Notes
    - mode has to be either 'joint' or 'cartesian'
    - data is a dataframe with the columns `step q1 q2 q3 q4 q5 q6` or `step x y z`
    - fk is a forward kinematic function to calculate cartesian points from joints. Only required if `mode=='joints'`
    - The `step` column is used to define correspondances between the set of points.

    Parameters
    ----------
    input_dict1: dict
        Input dict
    input_dict2: dict
        Input dict

    TODO:
    - Cartesian calculation from joints is not very intuitive for people to use.
    - Maybe extract cartesian should be embedded in the fk_function
    """

    input_dict1: dict
    input_dict2: dict

    def __post_init__(self):
        assert self.input_dict1["mode"] in ["joint", "cartesian"]
        assert self.input_dict2["mode"] in ["joint", "cartesian"]

        if self.input_dict1["mode"] == "joint":
            assert "fk" in self.input_dict1, "No fk function for input_data1"
            self.set1_cp_df = self.calculate_cartesian(
                self.input_dict1["data"], self.input_dict1["fk"]
            )
        else:
            self.set1_cp_df = self.input_dict1["data"]

        if self.input_dict2["mode"] == "joint":
            assert "fk" in self.input_dict2, "No fk function for input_data2"
            self.set2_cp_df = self.calculate_cartesian(
                self.input_dict2["data"], self.input_dict2["fk"]
            )
        else:
            self.set2_cp_df = self.input_dict2["data"]

    def calculate_fre(self):
        data = pd.merge(
            self.set1_cp_df, self.set2_cp_df, how="inner", on="step", suffixes=("_1", "_2")
        )

        A = data[["x_1", "y_1", "z_1"]].to_numpy().T
        B = data[["x_2", "y_2", "z_2"]].to_numpy().T

        rig_T = Frame.find_transformation_direct(A, B)
        B_est = rig_T @ A
        errors = B_est - B

        error_list = np.linalg.norm(errors, axis=0)

        return error_list

    def calculate_cartesian(self, data_df, fk_func):
        steps = data_df["step"].to_numpy().reshape((-1, 1))
        data_cp: spatialmath.pose3d.SE3
        data_cp = fk_func(data_df[["q1", "q2", "q3", "q4", "q5", "q6"]].to_numpy())
        data_cp = self.extract_cartesian_xyz(data_cp.data)
        data_cp = np.hstack((steps, data_cp))

        return pd.DataFrame(data_cp, columns=["step", "x", "y", "z"])

    def extract_cartesian_xyz(self, cartesian_t: np.ndarray):
        position_list = []
        for i in range(len(cartesian_t)):
            position_list.append(cartesian_t[i][:3, 3])

        return np.array(position_list)


import eventlet
from eventlet import greenpool

from oslo_config import cfg
from oslo_log import log as logging
from oslo_service import loopingcall
from oslo_utils import excutils
from oslo_utils import strutils
from oslo_utils import timeutils
from taskflow.patterns import linear_flow
from taskflow import retry

import masakari.conf
from masakari.engine.drivers.taskflow import base
from masakari import exception
from masakari import objects
from masakari.objects import fields
from masakari import utils

CONF = masakari.conf.CONF
LOG = logging.getLogger(__name__)
ACTION = 'instance:evacuate'
# Instance power_state
SHUTDOWN = 4
TASKFLOW_CONF = cfg.CONF.taskflow_driver_recovery_flows

class DisableComputeServiceTask(base.MasakariTask):
    def __init__(self, context, novaclient, **kwargs):
        kwargs['requires'] = ["host_name"]
        super(DisableComputeServiceTask, self).__init__(context, novaclient,
                                                        **kwargs)

    def execute(self, host_name):
        msg = "Disabling compute service on host: '%s'" % host_name
        self.update_details(msg)
        self.novaclient.enable_disable_service(self.context, host_name,
            reason=CONF.host_failure.service_disable_reason)
        # Sleep until nova-compute service is marked as disabled.
        log_msg = ("Sleeping %(wait)s sec before starting recovery "
               "thread until nova recognizes the node down.")
        LOG.info(log_msg, {'wait': CONF.wait_period_after_service_update})
        eventlet.sleep(CONF.wait_period_after_service_update)
        msg = "Disabled compute service on host: '%s'" % host_name
        self.update_details(msg, 1.0)

class PrepareHAEnabledInstancesTask(base.MasakariTask):
    """Get all HA_Enabled instances."""

    def __init__(self, context, novaclient, **kwargs):
        kwargs['requires'] = ["host_name", "notification_uuid"]
        super(PrepareHAEnabledInstancesTask, self).__init__(context,
                                                            novaclient,
                                                            **kwargs)

    def execute(self, host_name, notification_uuid):
        def _filter_instances(instance_list):
            ha_enabled_instances = []
            non_ha_enabled_instances = []

            ha_enabled_key = CONF.host_failure.ha_enabled_instance_metadata_key

            for instance in instance_list:
                is_instance_ha_enabled = strutils.bool_from_string(
                    instance.metadata.get(ha_enabled_key, False))
                if CONF.host_failure.ignore_instances_in_error_state and (
                        getattr(instance, "OS-EXT-STS:vm_state") == "error"):
                    if is_instance_ha_enabled:
                        msg = ("Ignoring recovery of HA_Enabled instance "
                               "'%(instance_id)s' as it is in 'error' state."
                               ) % {'instance_id': instance.id}
                        LOG.info(msg)
                        self.update_details(msg, 0.4)
                    continue

                if is_instance_ha_enabled:
                    ha_enabled_instances.append(instance)
                else:
                    non_ha_enabled_instances.append(instance)

            msg = "Total HA Enabled instances count: '%d'" % len(
                ha_enabled_instances)
            self.update_details(msg, 0.6)

            if CONF.host_failure.evacuate_all_instances:
                msg = ("Total Non-HA Enabled instances count: '%d'" % len(
                    non_ha_enabled_instances))
                self.update_details(msg, 0.7)

                ha_enabled_instances.extend(non_ha_enabled_instances)

                msg = ("All instances (HA Enabled/Non-HA Enabled) should be "
                       "considered for evacuation. Total count is: '%d'") % (
                    len(ha_enabled_instances))
                self.update_details(msg, 0.8)

            return ha_enabled_instances

        msg = "Preparing instances for evacuation"
        self.update_details(msg)

        instance_list = self.novaclient.get_servers(self.context, host_name)

        msg = ("Total instances running on failed host '%(host_name)s' is "
               "%(instance_list)d") % {'host_name': host_name,
                                       'instance_list': len(instance_list)}
        self.update_details(msg, 0.3)

        instance_list = _filter_instances(instance_list)

        if not instance_list:
            msg = ("Skipped host '%s' recovery as no instances needs to be "
                   "evacuated" % host_name)
            self.update_details(msg, 1.0)
            LOG.info(msg)
            raise exception.SkipHostRecoveryException(message=msg)

        # persist vm moves
        for instance in instance_list:
            vmove = objects.VMove(context=self.context)
            vmove.instance_uuid = instance.id
            vmove.instance_name = instance.name
            vmove.notification_uuid = notification_uuid
            vmove.source_host = host_name
            vmove.status = fields.VMoveStatus.PENDING
            vmove.type = fields.VMoveType.EVACUATION
            vmove.create()

        # List of instance UUID
        instance_list = [instance.id for instance in instance_list]

        msg = "Instances to be evacuated are: '%s'" % ','.join(instance_list)
        self.update_details(msg, 1.0)

class EvacuateInstancesTask(base.MasakariTask):

    def __init__(self, context, novaclient, **kwargs):
        kwargs['requires'] = ["host_name", "notification_uuid"]
        self.update_host_method = kwargs['update_host_method']
        super(EvacuateInstancesTask, self).__init__(context, novaclient,
                                                    **kwargs)

    def _get_state_and_host_of_instance(self, context, instance):
        new_instance = self.novaclient.get_server(context, instance.id)
        instance_host = getattr(new_instance,
                                "OS-EXT-SRV-ATTR:hypervisor_hostname")
        old_vm_state = getattr(instance, "OS-EXT-STS:vm_state")
        new_vm_state = getattr(new_instance, "OS-EXT-STS:vm_state")

        return (old_vm_state, new_vm_state, instance_host)

    def _stop_after_evacuation(self, context, instance):
        def _wait_for_stop_confirmation():
            old_vm_state, new_vm_state, instance_host = (
                self._get_state_and_host_of_instance(context, instance))

            if new_vm_state == 'stopped':
                raise loopingcall.LoopingCallDone()

        try:
            # confirm instance is stopped after recovery
            self.novaclient.stop_server(context, instance.id)
            timer = loopingcall.FixedIntervalWithTimeoutLoopingCall(
                _wait_for_stop_confirmation)
            timer.start(interval=CONF.verify_interval,
                        timeout=CONF.wait_period_after_power_off).wait()
        except loopingcall.LoopingCallTimeOut:
            with excutils.save_and_reraise_exception():
                msg = ("Instance '%(uuid)s' is successfully evacuated but "
                       "timeout to stop.") % {'uuid': instance.id}
                LOG.warning(msg)
        finally:
            timer.stop()

    def _evacuate_and_confirm(self, context, vmove,
                              reserved_host=None):

        def _update_vmove(vmove, status=None, start_time=None,
                          end_time=None, dest_host=None,
                          message=None):
            if status:
                vmove.status = status
            if start_time:
                vmove.start_time = start_time
            if end_time:
                vmove.end_time = end_time
            if dest_host:
                vmove.dest_host = dest_host
            if message:
                vmove.message = message
            vmove.save()

        instance_uuid = vmove.instance_uuid
        instance = self.novaclient.get_server(context, instance_uuid)

        # Before locking the instance check whether it is already locked
        # by user, if yes don't lock the instance
        instance_already_locked = self.novaclient.get_server(
            context, instance.id).locked

        if not instance_already_locked:
            # lock the instance so that until evacuation and confirmation
            # is not complete, user won't be able to perform any actions
            # on the instance.
            self.novaclient.lock_server(context, instance.id)

        def _wait_for_evacuation_confirmation():
            old_vm_state, new_vm_state, instance_host = (
                self._get_state_and_host_of_instance(context, instance))

            if (new_vm_state == 'error' and
                    new_vm_state != old_vm_state):
                raise exception.InstanceEvacuateFailed(
                    instance_uuid=instance.id)

            if instance_host != vmove.source_host:
                if ((old_vm_state == 'error' and
                    new_vm_state == 'active') or
                        old_vm_state == new_vm_state):
                    raise loopingcall.LoopingCallDone()

        def _wait_for_evacuation():
            try:
                # add a timeout to the periodic call.
                timer = loopingcall.FixedIntervalWithTimeoutLoopingCall(
                    _wait_for_evacuation_confirmation)
                timer.start(interval=CONF.verify_interval,
                            timeout=CONF.wait_period_after_evacuation).wait()
            except loopingcall.LoopingCallTimeOut:
                with excutils.save_and_reraise_exception():
                    msg = ("Timeout for instance '%(uuid)s' evacuation."
                           % {'uuid': instance.id})
                    LOG.warning(msg)
            finally:
                # stop the periodic call, in case of exceptions or
                # Timeout.
                timer.stop()

        try:
            vm_state = getattr(instance, "OS-EXT-STS:vm_state")
            task_state = getattr(instance, "OS-EXT-STS:task_state")

            # Nova evacuates an instance only when vm_state is in active,
            # stopped or error state. If an instance is in other than active,
            # error and stopped vm_state, masakari resets the instance state
            # to *error* so that the instance can be evacuated.
            stop_instance = True
            if vm_state not in ['active', 'error', 'stopped']:
                self.novaclient.reset_instance_state(context, instance.id)
                instance = self.novaclient.get_server(context, instance.id)
                power_state = getattr(instance, "OS-EXT-STS:power_state")
                if vm_state == 'resized' and power_state != SHUTDOWN:
                    stop_instance = False

            elif vm_state == 'stopped' and task_state is None:
                # If vm_state is stopped and task_state is none, the instance
                # will be recovered with vm_state 'stopped'.
                # So it doesn't need to stop the instance after evacuation.
                stop_instance = False

            elif task_state is not None:
                # Nova fails evacuation when the instance's task_state is not
                # none. In this case, masakari resets the instance's vm_state
                # to 'error' and task_state to none.
                self.novaclient.reset_instance_state(context, instance.id)
                instance = self.novaclient.get_server(context, instance.id)
                if vm_state == 'active':
                    stop_instance = False

            # start to evacuate the instance
            _update_vmove(
                vmove,
                status=fields.VMoveStatus.ONGOING,
                start_time=timeutils.utcnow())
            self.novaclient.evacuate_instance(context, instance.id,
                                              target=reserved_host)

            _wait_for_evacuation()

            if vm_state != 'active':
                if stop_instance:
                    self._stop_after_evacuation(self.context, instance)
                    # If the instance was in 'error' state before failure
                    # it should be set to 'error' after recovery.
                    if vm_state == 'error':
                        self.novaclient.reset_instance_state(
                            context, instance.id)

            instance = self.novaclient.get_server(context, instance_uuid)
            dest_host = getattr(
                instance, "OS-EXT-SRV-ATTR:hypervisor_hostname")
            _update_vmove(
                vmove,
                status=fields.VMoveStatus.SUCCEEDED,
                dest_host=dest_host)
        except loopingcall.LoopingCallTimeOut:
            # Instance is not stop in the expected time_limit.
            msg = "Failed reason: timeout."
            _update_vmove(
                vmove,
                status=fields.VMoveStatus.FAILED,
                message=msg)
        except Exception as e:
            # Exception is raised while resetting instance state or
            # evacuating the instance itself.
            LOG.warning(str(e))
            _update_vmove(
                vmove,
                status=fields.VMoveStatus.FAILED,
                message=str(e))
        finally:
            _update_vmove(vmove, end_time=timeutils.utcnow())
            if not instance_already_locked:
                # Unlock the server after evacuation and confirmation
                self.novaclient.unlock_server(context, instance.id)

    def execute(self, host_name, notification_uuid, reserved_host=None):
        all_vmoves = objects.VMoveList.get_all_vmoves(
            self.context, notification_uuid, status=fields.VMoveStatus.PENDING)
        instance_list = [i.instance_uuid for i in all_vmoves]
        msg = ("Start evacuation of instances from failed host '%(host_name)s'"
               ", instance uuids are: '%(instance_list)s'") % {
            'host_name': host_name, 'instance_list': ','.join(instance_list)}
        self.update_details(msg)

        def _do_evacuate(context, host_name,
                         reserved_host=None):
            if reserved_host:
                msg = "Enabling reserved host: '%s'" % reserved_host
                self.update_details(msg, 0.1)
                if CONF.host_failure.add_reserved_host_to_aggregate:
                    # Assign reserved_host to an aggregate to which the failed
                    # compute host belongs to.
                    aggregates = self.novaclient.get_aggregate_list(context)
                    for aggregate in aggregates:
                        if host_name in aggregate.hosts:
                            try:
                                msg = ("Add host %(reserved_host)s to "
                                       "aggregate %(aggregate)s") % {
                                    'reserved_host': reserved_host,
                                    'aggregate': aggregate.name}
                                self.update_details(msg, 0.2)

                                self.novaclient.add_host_to_aggregate(
                                    context, reserved_host, aggregate)
                                msg = ("Added host %(reserved_host)s to "
                                       "aggregate %(aggregate)s") % {
                                    'reserved_host': reserved_host,
                                    'aggregate': aggregate.name}
                                self.update_details(msg, 0.3)
                            except exception.Conflict:
                                msg = ("Host '%(reserved_host)s' already has "
                                       "been added to aggregate "
                                       "'%(aggregate)s'.") % {
                                    'reserved_host': reserved_host,
                                    'aggregate': aggregate.name}
                                self.update_details(msg, 1.0)
                                LOG.info(msg)

                self.novaclient.enable_disable_service(
                    context, reserved_host, enable=True)

                # Set reserved property of reserved_host to False
                self.update_host_method(context, reserved_host)

            thread_pool = greenpool.GreenPool(
                CONF.host_failure_recovery_threads)

            nonlocal all_vmoves

            for vmove in all_vmoves:
                msg = ("Evacuation of instance started: '%s'"
                       % vmove.instance_uuid)
                self.update_details(msg, 0.5)
                thread_pool.spawn_n(self._evacuate_and_confirm, self.context,
                                    vmove, reserved_host)
            thread_pool.waitall()

            all_vmoves = objects.VMoveList.get_all_vmoves(
                self.context, notification_uuid)

            succeeded_vmoves = [i.instance_uuid for i in all_vmoves
                    if i.status == fields.VMoveStatus.SUCCEEDED]
            if succeeded_vmoves:
                succeeded_vmoves.sort()
                msg = ("Successfully evacuate instances '%(instance_list)s' "
                       "from host '%(host_name)s'") % {
                    'instance_list': ','.join(succeeded_vmoves),
                    'host_name': host_name}
                self.update_details(msg, 0.7)

            failed_vmoves = [i.instance_uuid for i in
                    all_vmoves if i.status == fields.VMoveStatus.FAILED]
            if failed_vmoves:
                msg = ("Failed to evacuate instances "
                       "'%(instance_list)s' from host "
                       "'%(host_name)s'") % {
                    'instance_list': ','.join(failed_vmoves),
                    'host_name': host_name}
                self.update_details(msg, 0.7)
                raise exception.HostRecoveryFailureException(
                    message=msg)

            msg = "Evacuation process completed!"
            self.update_details(msg, 1.0)

        lock_name = reserved_host if reserved_host else None

        @utils.synchronized(lock_name)
        def do_evacuate_with_reserved_host(context, host_name,
                notification_uuid, reserved_host):
            _do_evacuate(context, host_name,
                         reserved_host=reserved_host)

        if lock_name:
            do_evacuate_with_reserved_host(self.context, host_name,
                                           notification_uuid,
                                           reserved_host)
        else:
            # No need to acquire lock on reserved_host when recovery_method is
            # 'auto' as the selection of compute host will be decided by nova.
            _do_evacuate(self.context, host_name)

def get_auto_flow(context, novaclient, process_what):
    """Constructs and returns the engine entrypoint flow.

    This flow will do the following:

    1. Disable compute service on source host
    2. Get all HA_Enabled instances.
    3. Evacuate all the HA_Enabled instances.
    4. Confirm evacuation of instances.
    """

    flow_name = ACTION.replace(":", "_") + "_engine"
    nested_flow = linear_flow.Flow(flow_name)

    task_dict = TASKFLOW_CONF.host_auto_failure_recovery_tasks

    auto_evacuate_flow_pre = linear_flow.Flow('pre_tasks')
    for plugin in base.get_recovery_flow(task_dict['pre'], context=context,
                                         novaclient=novaclient,
                                         update_host_method=None):
        auto_evacuate_flow_pre.add(plugin)

    auto_evacuate_flow_main = linear_flow.Flow('main_tasks')
    for plugin in base.get_recovery_flow(task_dict['main'], context=context,
                                         novaclient=novaclient,
                                         update_host_method=None):
        auto_evacuate_flow_main.add(plugin)

    auto_evacuate_flow_post = linear_flow.Flow('post_tasks')
    for plugin in base.get_recovery_flow(task_dict['post'], context=context,
                                         novaclient=novaclient,
                                         update_host_method=None):
        auto_evacuate_flow_post.add(plugin)

    nested_flow.add(auto_evacuate_flow_pre)
    nested_flow.add(auto_evacuate_flow_main)
    nested_flow.add(auto_evacuate_flow_post)

    return base.load_taskflow_into_engine(ACTION, nested_flow,
                                          process_what)

def get_rh_flow(context, novaclient, process_what, **kwargs):
    """Constructs and returns the engine entrypoint flow.

    This flow will do the following:

    1. Disable compute service on source host
    2. Get all HA_Enabled instances.
    3. Evacuate all the HA_Enabled instances using reserved_host.
    4. Confirm evacuation of instances.
    """
    flow_name = ACTION.replace(":", "_") + "_engine"
    nested_flow = linear_flow.Flow(flow_name)

    task_dict = TASKFLOW_CONF.host_rh_failure_recovery_tasks

    rh_evacuate_flow_pre = linear_flow.Flow('pre_tasks')
    for plugin in base.get_recovery_flow(
            task_dict['pre'],
            context=context, novaclient=novaclient, **kwargs):
        rh_evacuate_flow_pre.add(plugin)

    rh_evacuate_flow_main = linear_flow.Flow(
        "retry_%s" % flow_name, retry=retry.ParameterizedForEach(
            rebind=['reserved_host_list'], provides='reserved_host'))

    for plugin in base.get_recovery_flow(
            task_dict['main'],
            context=context, novaclient=novaclient, **kwargs):
        rh_evacuate_flow_main.add(plugin)

    rh_evacuate_flow_post = linear_flow.Flow('post_tasks')
    for plugin in base.get_recovery_flow(
            task_dict['post'],
            context=context, novaclient=novaclient, **kwargs):
        rh_evacuate_flow_post.add(plugin)

    nested_flow.add(rh_evacuate_flow_pre)
    nested_flow.add(rh_evacuate_flow_main)
    nested_flow.add(rh_evacuate_flow_post)

    return base.load_taskflow_into_engine(ACTION, nested_flow,
                                          process_what)