# coding=UTF-8
import cv2
import numpy as np

from basic_function import show_img, get_M_Minv, draw_area, draw_demo

def detect_line(img_input, steer, memory={"left_fit":[0.0, 0.0, 0.0],"right_fit":[0.0, 0.0, 0.0],"left_x":0.0,"right_x":0.0}, debug=False):
    """
    Main Function

    Parameters:
    img: original image
    steer: real steer of the vehicle to display
    memory: memory of the line-fit parameters

    Return:
    img_result: demo image
    distance_from_center:  positive--Right   negetive--Left   None--not sure
    curvature
    memory: memory of the line-fit parameters
    img_area: image for lane area
    """
    # Pre-process
    img = img_input[1000:, :]
    img_line, img_bin, img_canny = pre_process(img, debug)

    # Get warp_line
    M, Minv = get_M_Minv()
    img_line_warp = cv2.warpPerspective(img_line, M, img.shape[1::-1], flags=cv2.INTER_LINEAR)
    img_line_warp[:, :1200] = 0
    img_line_warp[:, 2000:] = 0
    kernel_ed = cv2.getStructuringElement(cv2.MORPH_RECT,(11,11))
    img_line_warp= cv2.dilate(img_line_warp,kernel_ed,2)
    if debug: show_img('warp', img_line_warp)

    # Detect line & Calculate the value
    left_fit, right_fit, left_x, right_x, img_bev_result = find_line(img_line_warp, memory, debug)
    curvature, distance_from_center = calculate_curv_and_pos(img_line_warp, left_fit, right_fit)

    # # Draw results on the image
    img_area = draw_area(img, img_bev_result, Minv, left_fit, right_fit)
    img_result = img_result = cv2.addWeighted(img, 1, img_area, 0.3, 0)
    img_result = np.vstack([img_input[0:1001, :], img_result])
    img_area = np.vstack([np.zeros_like(img_input)[0:1001, :], img_area])
    img_result, distance_from_center = draw_demo(img_result, img_bin, img_canny, img_line, img_line_warp, img_bev_result, curvature, distance_from_center, steer)

    memory = {"left_fit":left_fit, "right_fit":right_fit, "left_x":left_x, "right_x":right_x}

    return img_result, distance_from_center, curvature,  memory, img_area

def pre_process(img, debug=False):
    """
    Image Preprocessing

    Parameters:
    img: original image

    Return:
    img_line
    img_bin
    img_canny
    """
    # BGR2GRAY
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Bin
    _, img_bin = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    if debug: show_img("BIN", img_bin)

    # Canny
    img_blur = cv2.GaussianBlur(img_gray,(5,5),0) # 5
    if debug: show_img("blur", img_blur)
    img_canny = cv2.Canny(img_blur, 10, 60, apertureSize=3) #10 50

    if debug: show_img("Canny", img_canny)

    # Bin & Canny
    img_line = cv2.bitwise_and(img_bin, img_canny)

    kernel_ed = cv2.getStructuringElement(cv2.MORPH_RECT,(11,11))
    img_line = cv2.dilate(img_line,kernel_ed,2)
    img_line = cv2.erode(img_line,kernel_ed,2)

    if debug: show_img("Line", img_line)

    return img_line, img_bin, img_canny

def find_line(img, memory, debug=False):
    """
    Detect the lane using Sliding Windows Methods

    Parameters:
    img: warp_line (Bin)
    memory: memory of the line-fit parameters

    Return:
    left_fit, right_fit
    format [a,b,c]
    y = a*x^2 + b*x + c
    """
    # 在x方向上统计y方向上的像素和，推断道路线可能存在的x位置，进而确定左右侧道路线的起始点
    histogram = np.sum(img[img.shape[0]//2:,:], axis=0)

    midpoint = np.int(histogram.shape[0]/2) + 300
    leftx_base = np.argmax(histogram[:midpoint])
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint
    # if memory:
    #     if abs(np.argsort(histogram[:midpoint])[-1]-memory['left_x']) < 100 or abs(np.argsort(histogram[:midpoint])[-1]-memory['left_x']) > 400 or abs(np.argsort(histogram[:midpoint])[-2]-memory['left_x']) < 100:
    #         leftx_base = np.argsort(histogram[:midpoint])[-1]
    #     else:
    #         leftx_base = memory['left_x']
    #         print ("l", min(abs(np.argsort(histogram[:midpoint])[-2]-memory['left_x']), abs(np.argsort(histogram[:midpoint])[-1]-memory['left_x'])), np.argsort(histogram[:midpoint])[-2], memory['left_x'])

    #     if abs(np.argsort(histogram[midpoint:])[-1]+midpoint-memory['right_x']) < 100 or abs(np.argsort(histogram[midpoint:])[-1]+midpoint-memory['right_x'])>400 or abs(np.argsort(histogram[midpoint:])[-2]+midpoint-memory['right_x']) < 100:
    #         rightx_base = np.argsort(histogram[midpoint:])[-1] + midpoint
    #     else:
    #         rightx_base = memory['right_x']
    #         print ("r", min(abs(np.argsort(histogram[midpoint:])[-2]+midpoint-memory['right_x']),abs(np.argsort(histogram[midpoint:])[-1]+midpoint-memory['right_x'])), np.argsort(histogram[midpoint:])[-2]+midpoint, memory['right_x'])

    #     print ()
    # else:
    #     leftx_base = np.argmax(histogram[:midpoint])
    #     rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # Initialize the windows value
    windows_num = 10 # 9
    windows_h = np.int(img.shape[0]/windows_num)
    windows_w = 100
    # Current positions to be updated for each window
    leftx_current = leftx_base
    rightx_current = rightx_base
    # Set minimum number of pixels found to recenter window
    minpix = 50 # 50

    # Identify the x and y positions of all nonzero pixels in the image
    nonzero = img.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])

    # Create empty lists to receive left and right lane pixel indices
    left_lane_inds = []
    right_lane_inds = []

    # GRAY2BGR
    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    # Step through the windows one by one
    for i in range(windows_num):
        # Identify window boundaries in x and y (and right and left)
        win_y_low = img.shape[0] - (i+1)*windows_h
        win_y_high = img.shape[0] - i*windows_h
        win_xleft_low = leftx_current - windows_w//2
        win_xleft_high = leftx_current + windows_w//2
        win_xright_low = rightx_current - windows_w//2
        win_xright_high = rightx_current + windows_w//2

        # Draw the window
        cv2.rectangle(img,(win_xleft_low,win_y_low),(win_xleft_high,win_y_high), (0,255,0), 2)
        cv2.rectangle(img,(win_xright_low,win_y_low),(win_xright_high,win_y_high), (0,255,0), 2)

        # Identify the nonzero pixels in x and y within the window
        good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & (nonzerox >= win_xleft_low) & (nonzerox < win_xleft_high)).nonzero()[0]
        good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & (nonzerox >= win_xright_low) & (nonzerox < win_xright_high)).nonzero()[0]
        # Append these indices to the lists
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)

        # If you found > minpix pixels, recenter next window on their mean position
        if len(good_left_inds) > minpix:
            leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
        if len(good_right_inds) > minpix:
            rightx_current = np.int(np.mean(nonzerox[good_right_inds]))

    # Concatenate the arrays of indices
    left_lane_inds = np.concatenate(left_lane_inds)
    right_lane_inds = np.concatenate(right_lane_inds)

    # Extract left and right line pixel positions
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    # Draw the left & right line
    for i in range(len(leftx)): img[lefty[i]][leftx[i]] = [0,0,255]
    for i in range(len(rightx)): img[righty[i]][rightx[i]] = [255,0,0]

    # Fit a second order polynomial to each
    # left_fit = np.polyfit(lefty, leftx, 2)
    # right_fit = np.polyfit(righty, rightx, 2)
    try:
        left_fit = np.polyfit(lefty, leftx, 2)
    except:
        left_fit = memory['left_fit']

    try:
        right_fit = np.polyfit(righty, rightx, 2)
    except:
        right_fit = memory['right_fit']

    return left_fit, right_fit, leftx_base, rightx_base, img

def calculate_curv_and_pos(img_line, left_fit, right_fit):
    """
    Calculate the curvature & distance from the center

    Parameters:
    img_line: warp_line (Bin)
    left_fit: [a,b,c]
    right_fit: [a,b,c]

    Return:
    curvature, distance_from_center
    """
    # Define y-value where we want radius of curvature
    ploty = np.linspace(0, img_line.shape[0]-1, img_line.shape[0] )
    leftx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
    rightx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]

    # Define conversions in x and y from pixels space to meters
    ym_per_pix = 20/1000
    xm_per_pix = 3.7/2030
    y_eval = np.max(ploty)

    # Fit new polynomials to x,y in world space
    left_fit_cr = np.polyfit(ploty*ym_per_pix, leftx*xm_per_pix, 2)
    right_fit_cr = np.polyfit(ploty*ym_per_pix, rightx*xm_per_pix, 2)

    # Calculate the new radii of curvature
    left_curverad = ((1 + (2*left_fit_cr[0]*y_eval*ym_per_pix + left_fit_cr[1])**2)**1.5) / np.absolute(2*left_fit_cr[0])
    right_curverad = ((1 + (2*right_fit_cr[0]*y_eval*ym_per_pix + right_fit_cr[1])**2)**1.5) / np.absolute(2*right_fit_cr[0])
    curvature = ((left_curverad + right_curverad) / 2)

    #print(curvature)
    lane_width = np.absolute(leftx[60] - rightx[60])
    lane_xm_per_pix = 3.7 / lane_width
    veh_pos = (((leftx[60] + rightx[60]) * lane_xm_per_pix) / 2.)
    cen_pos = ((img_line.shape[1] * lane_xm_per_pix) / 2.)
    distance_from_center = cen_pos - veh_pos + 1.8

    return curvature, distance_from_center


import datetime
import numpy as np
import pandas as pd
import json
import dask.bag as daskBag
import dask.dataframe as dd
from dask.threaded import get
from dask.delayed import delayed
from dask.optimization import cull,inline,inline_functions,fuse
from yuntu.core.db.utils import timeAsCat
from yuntu.core.audio.base import Audio

def writeParquet(node,path):
    print("Writing parquet to "+path)
    return node.to_parquet(path,compression="GZIP")

def readParquet(path,npartitions):
    return dd.read_parquet(path)

def projectMetadata(obj,fields):
    newObj = {}
    newObj["orid"] = obj["orid"]
    newObj["media_info"] = obj["media_info"]
    newObj["fileDuration"] = obj["media_info"]["duration"]

    for key in fields:
        key_split = key.split(".")
        if len(key_split) > 1:

            val = obj["metadata"][key_split[0]]
            for i in range(1,len(key_split)):
                val = val[key_split[i]]

            newObj[key] = val
        else:
            if key in obj["metadata"]:
                newObj[key] = obj["metadata"][key]

    return newObj

def doTransform(e,eTransform,transformDict):
    if eTransform is not None:
        kwargs = {}
        if "kwargs" in transformDict:
            if transformDict["kwargs"] is not None:
                kwargs = transformDict["kwargs"]

        return eTransform(e,**kwargs)
    else:
        return e

def sliceEnergy(spec,eCols,fCuts,duration,start,stop,unitSize,eTransform,config):
    results = {}
    try:
        fbins,tbins = spec.shape
        if stop - start > config["tStep"]:
            e = spec
            nbins = tbins
        else:
            startBin = min(int(round((float(tbins)/duration)*start)),tbins-1)
            stopBin = min(startBin+unitSize,tbins-1)
            nbins = stopBin-startBin
            e = spec[:,startBin:stopBin]

        if nbins < float(unitSize)/2:
            raise ValueError("Generated empty slice!")

        results["chunkTbins"] = nbins
        results["chunkWeight"] = float(nbins)/unitSize

        if eTransform["full_spec"] != "pass":
            tConfig = config["transformations"]["full_spec"].copy()
            if "kwargs" in tConfig:
                #Use available context variables
                for kname in tConfig["kwargs"]:
                    if tConfig["kwargs"][kname] == "$fCuts":
                        tConfig["kwargs"][kname] = fCuts

            global_results =  doTransform(e,eTransform["full_spec"],tConfig)
            #Ensure transformation returns a dictionary with correct keys.
            for i in range(len(eCols)):
                results[eCols[i]] = global_results[eCols[i]]

            return results

        if eTransform["spec"] != "pass":
            e = doTransform(e,eTransform["spec"],config["transformations"]["spec"])
        #e = np.sum(e,axis=1)

        for i in range(len(eCols)):
            cut = fCuts[i]
            if cut[0] != cut[1]:
                eCut = e[cut[0]:cut[1],:]
                eCut = doTransform(eCut,eTransform["aggr"],config["transformations"]["aggr"])
            else:
                eCut = None

            results[eCols[i]] = eCut
    except:
        for col in ["chunkTbins","chunkWeight"]+eCols:
            results[col] = 0

    return results

def splitTime(obj,eCols,eTransform,config):
    prevCols = [key for key in list(obj.keys()) if key != "media_info"]
    tStep = config["tStep"]
    media_info = obj["media_info"]
    fileDuration = media_info["duration"]
    sTransform = eTransform["signal"]

    au = Audio(config=media_info,fromConfig=True)
    au.unsetMask()

    sr = media_info["samplerate"]
    if config["readSr"] is not None:
        if config["readSr"] != media_info["samplerate"]:
            sr = config["readSr"]
            au.setReadSr(config["readSr"])

    unitSize = int(round(float(sr*tStep-config["n_fft"])/config["hop_length"])+1)
    nsteps = int(round(fileDuration/tStep))

    chanList = [0]

    if config["channels"] is None:
        chanList = [None]
    elif config["channels"] == "__all__":
        if media_info["nchannels"] > 1:
            chanList = range(media_info["nchannels"])
    elif isinstance(config["channels"],list):
        chanList = config["channels"]
        if len(chanList) > media_info["nchannels"]:
            raise ValueError("Too many channels")
    else:
        raise ValueError("'channel' param should be list,str or None")

    preProcess = None
    if sTransform != "pass":
        kwargs = {}
        if "kwargs" in config["transformations"]["signal"]:
            if config["transformations"]["signal"]["kwargs"] is not None:
                kwargs = config["transformations"]["signal"]["kwargs"]

        preProcess = {"transform":sTransform,"kwargs":kwargs}

    out = []
    for chan in chanList:
        #try:
        spec, freqs = au.getSpec(chan,config["n_fft"],config["hop_length"],preProcess=preProcess)

        if chan is None:
            chan = 99

        topFreq = np.amax(freqs)
        fbins = freqs.size
        maxFreqIdx = fbins
        minFreqIdx = 0
        freqRange = topFreq

        if config["fLimits"] is not None:
            freqRange = config["fLimits"][1]-config["fLimits"][0]
            maxFreqIdx = min(int(round(fbins*config["fLimits"][1]/freqRange)),fbins)
            minFreqIdx = min(int(round(fbins*config["fLimits"][0]/freqRange)),fbins)

            if maxFreqIdx <= minFreqIdx:
                raise ValueError("Wrong frequency limits")

            spec = spec[minFreqIdx:maxFreqIdx,:]
        fStep = round(int(float(freqRange)/config["fBins"]))
        fSteps = [[i*fStep,(i+1)*fStep] for i in range(config["fBins"])]
        fCuts =  [[(np.abs(freqs - x[0])).argmin(),(np.abs(freqs - x[1])).argmin()] for x in fSteps]

        #except:
        #    print("Error during energy processing...")
        #    spec = None

        if spec is not None:
            for step in range(nsteps):
                raw = {}
                for key in prevCols:
                    raw[key] = obj[key]

                fileStart = step*tStep
                fileStop = min(fileStart+tStep,obj["fileDuration"])
                duration = fileStop-fileStart

                startSec = fileStart+obj["absStart"]
                stopSec = startSec+duration

                raw["channel"] = chan
                raw["chunkId"] = str(obj["orid"])+"-"+str(startSec)+"-"+str(stopSec)
                raw["chunkFileStart"] = fileStart
                raw["chunkFileStop"] = fileStop
                raw["chunkDuration"] = duration
                raw["chunkAbsStart"] = startSec
                raw["chunkAbsStop"] = stopSec

                eResults = sliceEnergy(spec,eCols,fCuts,fileDuration,fileStart,fileStop,unitSize,eTransform,config)
                for col in ["chunkTbins","chunkWeight"]+eCols:
                    raw[col] = eResults[col]
                out.append(raw)

    au.clearMedia()

    return out

def calendarize(row,config,fieldName,transformTz):
    row[fieldName] = timeAsCat(row["standardStart"],config["startDate"],config["timeZone"],config["timeFormat"],config["unit"],config["modulo"],transformTz)

    return row

def loadFragment(data,groupfields,config):
    columns = ["orid","media_info","fileDuration","standardStart","standardStop","absStart","absStop"]+groupfields

    return daskBag.from_sequence(data,npartitions=config["npartitions"]).map(projectMetadata,columns)

def makeSplits(fragment,basicmeta,splitmeta,groupmeta,eCols,eTransform,config):
    meta = basicmeta+splitmeta+groupmeta
    meta.append(("channel",np.dtype('int64')))

    for col in eCols:
        meta.append((col, np.dtype('float64')))

    dobjs = [delayed(splitTime)(row,eCols,eTransform,config) for row in fragment]

    return daskBag.from_delayed(dobjs).to_dataframe(meta=meta)

def makeCalendar(splits,config):
    meta = [x for x in zip(splits.columns.values,splits.dtypes.values)]+[("timeCell",np.dtype('int64'))]

    return splits.apply(calendarize,meta=meta,axis=1,config=config,fieldName="timeCell",transformTz=None)

def makeCronoCounts(splits):
    return splits.groupby(['timeCell'])["orid"].agg(['count'])

def makeCounts(splits):
    return len(splits)

def makeSample(splits,groupfields,config):
    if config["size"] == -1:
        return splits
    else:
        groupcols = ["timeCell"]
        if groupfields is not None:
            groupcols += groupfields

        meta = [x for x in zip(splits.columns.values,splits.dtypes.values)]

        return splits.groupby(groupcols).apply(lambda x: x.sample(frac=float(config["size"])/len(x),replace=config["replace"]) if len(x) > config["size"] else x,meta=meta).reset_index(drop=True)

def makeStats(statKey):
    def f(sample):
        return sample

    return f

def dGet(graph,opList,client):
    if client is not None:
        return client.get(graph,opList,sync=False)
    else:
        return get(graph,opList)

def linearizeOps(graph,opList):
    graph1, deps = cull(graph, opList)
    graph2 = inline(graph1, dependencies=deps)
    graph3 = inline_functions(graph2, opList, [len, str.split],dependencies=deps)
    graph4, deps = fuse(graph3)

    return graph4


import time, datetime, os
import threading
import numpy as np
import torch
import copy
import json
from typing import Union
from abc import ABC, abstractmethod
from rlib.utils.utils import fold_batch
from rlib.utils.VecEnv import BatchEnv, DummyBatchEnv
import torch
from torch.utils.tensorboard import SummaryWriter

class SyncMultiEnvTrainer(object):
    def __init__(self, envs: Union[BatchEnv, DummyBatchEnv], model:torch.nn.Module, val_envs: Union[list, BatchEnv, DummyBatchEnv], train_mode='nstep', return_type='nstep', log_dir='logs/', model_dir='models/', total_steps=50e6, nsteps=5, gamma=0.99, lambda_=0.95,
                     validate_freq=1e6, save_freq=0, render_freq=0, update_target_freq=0, num_val_episodes=50, max_val_steps=10000, log_scalars=True):
        '''
            A synchronous multiple env training framework for pytorch

            Args:
                envs - BatchEnv | DummyBatchEnv: multiple synchronous training environments
                model - reinforcement learning model
                log_dir, log directory string for location of directory to log scalars log_dir='logs/', model_dir='models/',
                val_envs - use your own discretion to choose which validation mode you wan't, recommended BatchEnv or list for Atari and DummyBatchEnv for Classic Control like envs
                    list: a list of envs for validation, uses threading to run environments asychronously
                    BatchEnv: uses multiprocessing to run validation envs sychronously in parallel
                    DummyBatchEnv: allows for sychronous env stepping without the overhead of multiprocessing, good for computationally cheap environments
                train_mode - 'nstep' or 'onestep' species whether training is done using multiple step TD learning or single step
                return_type - string to determine whether 'nstep', 'lambda' or 'GAE' returns are to be used
                total_steps - number of Total training steps across all environements
                nsteps - number of steps TD error is caluclated over
                validate_freq - number of steps across all environements before performing validating, 0 for no validation
                save_freq - number of steps across all environements before saving model, 0 for no saving
                render_freq - multiple of validate_freq before rendering (i.e. render every X validations), 0 for no rendering
                update_target_freq - number of steps across all environements before updating target model, 0 for no updating
                num_val_episodes - number of episodes to average over when validating
                max_val_steps - maximum number of steps for each validation episode (prevents infinite loops)
                log_scalars - boolean flag whether to log tensorboard scalars to log_dir
        '''
        self.env = envs
        if isinstance(envs, list):
            self.validate_func = self.validate_async
        else:
            self.validate_func = self.validate_sync
        if train_mode not in ['nstep', 'onestep']:
            raise ValueError('train_mode %s is not a valid argument. Valid arguments are ... %s, %s' %(train_mode,'nstep','onestep'))
        assert num_val_episodes >= len(val_envs), 'number of validation epsiodes {} must be greater than or equal to the number of validation envs {}'.format(num_val_episodes, len(val_envs))
        if return_type not in ['nstep', 'lambda', 'GAE']:
            raise ValueError('return_type %s is not a valid argument. Valid arguments are ... %s, %s, %s' %(return_type, 'nstep', 'lambda', 'GAE'))
        self.train_mode = train_mode
        self.num_envs = len(envs)
        self.env_id = envs.spec.id
        self.val_envs = val_envs
        self.validate_rewards = []
        self.model = model

        self.total_steps = int(total_steps)
        self.nsteps = nsteps
        self.return_type = return_type
        self.gamma = gamma
        self.lambda_ = lambda_

        self.validate_freq = int(validate_freq)
        self.num_val_episodes = num_val_episodes
        self.val_steps = max_val_steps
        self.lock = threading.Lock()

        self.save_freq = int(save_freq)
        self.render_freq = render_freq
        self.target_freq = int(update_target_freq)
        self.s = 0 # number of saves made
        self.t = 1 # number of updates done
        self.log_scalars = log_scalars
        self.log_dir = log_dir
        self.model_dir = model_dir

        self.states = self.env.reset()

        if log_scalars:
            # Tensorboard Variables
            self.train_log_dir = self.log_dir  + '/train'
            self.train_writer = SummaryWriter(self.train_log_dir)

        if not os.path.exists(self.model_dir) and save_freq > 0:
            os.makedirs(self.model_dir)

    def __del__(self):
        self.env.close()

    def train(self):
        if self.train_mode == 'nstep':
            self._train_nstep()
        elif self.train_mode == 'onestep':
            self._train_onestep()
        else:
            raise ValueError('%s is not a valid training mode'%(self.train_mode))

    @abstractmethod
    def _train_nstep(self):
        '''
            template for multi-step training loop for synchronous training over multiple environments
        '''
        start = time.time()
        batch_size = self.num_envs * self.nsteps
        num_updates = self.total_steps // batch_size
        # main loop
        for t in range(self.t,num_updates+1):
            states, actions, rewards, dones, values, last_values = self.rollout()
            if self.return_type == 'nstep':
                R = self.nstep_return(rewards, last_values, dones, gamma=self.gamma)
            elif self.return_type == 'GAE':
                R = self.GAE(rewards, values, last_values, dones, gamma=self.gamma, lambda_=self.lambda_) + values
            elif self.return_type == 'lambda':
                R = self.lambda_return(rewards, values, last_values, dones, gamma=self.gamma, lambda_=self.lambda_, clip=False)
            # stack all states, actions and Rs from all workers into a single batch
            states, actions, R = fold_batch(states), fold_batch(actions), fold_batch(R)
            l = self.model.backprop(states, R, actions)

            if self.render_freq > 0 and t % ((self.validate_freq // batch_size) * self.render_freq) == 0:
                render = True
            else:
                render = False

            if self.validate_freq > 0 and t % (self.validate_freq // batch_size) == 0:
                self.validation_summary(t,l,start,render)
                start = time.time()

            if self.save_freq > 0 and  t % (self.save_freq // batch_size) == 0:
                self.s += 1
                self.save(self.s)
                print('saved model')

            if self.target_freq > 0 and t % (self.target_freq // batch_size) == 0: # update target network (for value based learning e.g. DQN)
                self.update_target()

            self.t +=1

    @abstractmethod
    def rollout(self):
        raise NotImplementedError(self, 'No rollout method found')

    def nstep_return(self, rewards, last_values, dones, gamma=0.99, clip=False):
        if clip:
            rewards = np.clip(rewards, -1, 1)

        T = len(rewards)

        # Calculate R for advantage A = R - V
        R = np.zeros_like(rewards)
        R[-1] = last_values * (1-dones[-1])

        for i in reversed(range(T-1)):
            # restart score if done as BatchEnv automatically resets after end of episode
            R[i] = rewards[i] + gamma * R[i+1] * (1-dones[i])

        return R

    def lambda_return(self, rewards, values, last_values, dones, gamma=0.99, lambda_=0.8, clip=False):
        if clip:
            rewards = np.clip(rewards, -1, 1)
        T = len(rewards)
        # Calculate eligibility trace R^lambda
        R = np.zeros_like(rewards)
        R[-1] =  last_values * (1-dones[-1])
        for t in reversed(range(T-1)):
            # restart score if done as BatchEnv automatically resets after end of episode
            R[t] = rewards[t] + gamma * (lambda_* R[t+1] + (1.0-lambda_) * values[t+1]) * (1-dones[t])

        return R

    def GAE(self, rewards, values, last_values, dones, gamma=0.99, lambda_=0.95, clip=False):
        if clip:
            rewards = np.clip(rewards, -1, 1)
        # Generalised Advantage Estimation
        Adv = np.zeros_like(rewards)
        Adv[-1] = rewards[-1] + gamma * last_values * (1-dones[-1]) - values[-1]
        T = len(rewards)
        for t in reversed(range(T-1)):
            delta = rewards[t] + gamma * values[t+1] * (1-dones[t]) - values[t]
            Adv[t] = delta + gamma * lambda_ * Adv[t+1] * (1-dones[t])

        return Adv

    def validation_summary(self, t, loss, start, render):
        batch_size = self.num_envs * self.nsteps
        tot_steps = t * batch_size
        time_taken = time.time() - start
        frames_per_update = (self.validate_freq // batch_size) * batch_size
        fps = frames_per_update / time_taken

        score = self.validate_func(render)
        print("update %i, validation score %f, total steps %i, loss %f, time taken for %i frames:%fs, fps %f \t\t\t" %(t,score,tot_steps,loss,frames_per_update,time_taken,fps))

        if self.log_scalars:
            self.train_writer.add_scalar('validation/score', score, tot_steps)
            self.train_writer.add_scalar('train/loss', loss, tot_steps)

    def save_model(self, s):
        model_loc = f'{self.model_dir}/{s}.pt'
        # default saving method is to save session
        torch.save(self.model.state_dict(), model_loc)

    def load_model(self, modelname, model_dir="models/"):
        filename = model_dir + modelname + '.pt'
        if os.path.exists(filename):
            self.model.load_state_dict(torch.load(filename))
            print("loaded:", filename)
        else:
            print(filename, " does not exist")

    def base_attr(self):
        attributes = {'train_mode':self.train_mode,
                'total_steps':self.total_steps,
                'nsteps':self.nsteps,
                'return_type':self.return_type,
                'gamma':self.gamma,
                'lambda_':self.lambda_,
                'validate_freq':self.validate_freq,
                'num_val_episodes':self.num_val_episodes,
                'save_freq':self.save_freq,
                'render_freq':self.render_freq,
                'model_dir':self.model_dir,
                'train_log_dir':self.train_log_dir,
                's':self.s,
                't':self.t}

        return attributes

    def local_attr(self, attr):
        # attr[variable] = z
        return attr

    def save(self, s):
        model_loc = str(self.model_dir + '/' + str(s) + '.trainer')
        file = open(model_loc, 'w+')
        attributes = self.base_attr()
        # add local variables to dict
        attributes = self.local_attr(attributes)
        json.dump(attributes, file)
        # save model
        self.save_model(s)
        file.close()

    def load(self, Class, model, model_checkpoint, envs, val_envs, filename, log_scalars=True, allow_gpu_growth=True, continue_train=True):
        with open(filename, 'r') as file:
            attrs = json.loads(file.read())
        s = attrs.pop('s')
        t = attrs.pop('t')
        time = attrs.pop('current_time')
        print(attrs)
        trainer = Class(envs=envs, model=model, val_envs=val_envs, log_scalars=log_scalars, gpu_growth=allow_gpu_growth, **attrs)
        if continue_train:
            trainer.s = s
            trainer.t = t
        self.load_model(model_checkpoint, trainer.model_dir)
        return trainer

    @abstractmethod
    def update_target(self):
        pass

    @abstractmethod
    def _train_onestep(self):
        ''' more efficient implementation of train_nstep when nsteps=1
        '''
        raise NotImplementedError(self, 'does not have an one-step training implementation')

    def save_hyperparameters(self, filename, **kwargs):
        handle = open(filename, "w")
        for key, value in kwargs.items():
            handle.write("{} = {}\n" .format(key, value))
        handle.close()

    def validate_async(self, render=False):
        num_val_envs = len(self.val_envs)
        num_val_eps = [self.num_val_episodes//num_val_envs for i in range(num_val_envs)]
        num_val_eps[-1] = num_val_eps[-1] + self.num_val_episodes % self.num_val_episodes//(num_val_envs)
        render_array = np.zeros((len(self.val_envs)))
        render_array[0] = render
        threads = [threading.Thread(daemon=True, target=self._validate_async, args=(self.val_envs[i], num_val_eps[i], self.val_steps, render_array[i])) for i in range(num_val_envs)]

        try:
            for thread in threads:
                thread.start()

            for thread in threads:
                thread.join()

        except KeyboardInterrupt:
            for thread in threads:
                thread.join()

        score = np.mean(self.validate_rewards)
        self.validate_rewards = []
        return score

    def _validate_async(self, env, num_ep, max_steps, render=False):
        'single env validation'
        for episode in range(num_ep):
            state = env.reset()
            episode_score = []
            for t in range(max_steps):
                action = self.get_action(state[np.newaxis])
                next_state, reward, done, info = env.step(action)
                state = next_state
                #print('state', state, 'action', action, 'reward', reward)

                episode_score.append(reward)

                if render:
                    with self.lock:
                        env.render()

                if done or t == max_steps -1:
                    tot_reward = np.sum(episode_score)
                    with self.lock:
                        self.validate_rewards.append(tot_reward)

                    break
        if render:
            with self.lock:
                env.close()

    def validate_sync(self, render=False):
        'batch env validation'
        episode_scores = []
        env = self.val_envs
        for episode in range(self.num_val_episodes//len(env)):
            states = env.reset()
            episode_score = []
            for t in range(self.val_steps):
                actions = self.get_action(states)
                next_states, rewards, dones, infos = env.step(actions)
                states = next_states
                #print('state', state, 'action', action, 'reward', reward)

                episode_score.append(rewards*(1-dones))

                if render:
                    with self.lock:
                        env.render()

                if dones.sum() == self.num_envs or t == self.val_steps -1:
                    tot_reward = np.sum(np.stack(episode_score), axis=0)
                    episode_scores.append(tot_reward)
                    break

        return np.mean(episode_scores)

    def get_action(self, state): # include small fn in order to reuse validate
        raise NotImplementedError('get_action method is required when using the default validation functions, check that this is implemented properly')

    def fold_batch(self, x):
        rows, cols = x.shape[0], x.shape[1]
        y = x.reshape(rows*cols,*x.shape[2:])
        return y

# class Runner(ABC):
#     def __init__(self,model,env,num_steps):
#         self.model = model
#         self.env = env
#         self.num_steps = num_steps
#         self.states = self.env.reset()

#     @abstractmethod
#     def run(self):
#         pass


#!/usr/bin/python3

"""Swing molecules based on normal coordinates."""

import os
import argparse
import shutil
from collections import defaultdict

import numpy as np
import matplotlib.pyplot as plt
from ase import Atoms
from ase.calculators.orca import ORCA
from cclib import ccopen
from scipy.constants import physical_constants
from scipy.optimize import root
from scipy.optimize import minimize
from scipy.optimize import minimize_scalar

try:
    from xtb.ase.calculator import XTB
except ModuleNotFoundError as e:
    print("WARNING: won't be able to use XTB")

hartree, _, _ = physical_constants["Hartree energy"]

smd2gbsa = defaultdict(
    lambda: "none",
    {
        "acetone": "acetone",
        "mecn": "acetonitrile",
        "acetonitrile": "acetonitrile",
        "benzene": "benzene",
        "dichloromethane": "ch2cl2",
        "chloroform": "chcl3",
        "carbon disulfide": "cs2",
        "dmf": "dmf",
        "n,n-dimethylformamide": "dmf",
        "dmso": "dmso",
        "dimethylsulfoxide": "dmso",
        "diethyl ether": "ether",
        "water": "water",
        "methanol": "methanol",
        "n-hexane": "n-hexan",
        "thf": "thf",
        "tetrahydrofuran": "thf",
        "toluene": "toluene",
    },
)

def main():
    """Run main procedure."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("logfile")
    # TODO(schneiderfelipe): set charge and multiplicity
    parser.add_argument(
        "-a",
        "--acc",
        help="accuracy for SCC calculation, lower is better",
        type=float,
        default=1.0,
    )
    parser.add_argument(
        "--iterations",
        help="number of iterations in SCC",
        type=int,
        default=250,
    )
    parser.add_argument(
        "--gfn", help="specify parametrisation of GFN-xTB", type=int
    )
    parser.add_argument(
        "--etemp", help="electronic temperature", type=float, default=300.0
    )
    parser.add_argument(
        "-s",
        "--solvent",
        help=("solvent (SMD/GBSA implicit solvation models)"),
        default="none",
    )
    parser.add_argument(
        "--do-not-cache-api",
        dest="cache_api",
        help="Do not reuse generate API objects (not recommended)",
        action="store_false",
    )

    parser.add_argument(
        "--pm3", help="use PM3", action="store_true",
    )
    parser.add_argument(
        "--b97-3c", help="use B97-3c", action="store_true",
    )
    parser.add_argument(
        "--minimize", action="store_true",
    )
    parser.add_argument(
        "--transition-state", action="store_true",
    )
    parser.add_argument("--max-omega", type=float, default=1.0)
    parser.add_argument("--tol", type=float, default=1e-3)
    parser.add_argument("--nprocs", type=int, default=4)
    args = parser.parse_args()
    print(args)

    data = ccopen(args.logfile).parse()
    initial_positions = data.atomcoords[-1]
    atoms = Atoms(numbers=data.atomnos, positions=initial_positions)

    if args.gfn:
        method = f"GFN{args.gfn}-xTB"
        solvent = smd2gbsa[args.solvent.lower()]

        calc = XTB(
            method=method,
            accuracy=args.acc,
            electronic_temperature=args.etemp,
            max_iterations=args.iterations,
            solvent=solvent,
            cache_api=args.cache_api,
        )
    else:

        if args.b97_3c:
            method = "B97-3c D3BJ def2-SV(P)"
        elif args.pm3:
            method = "PM3"
        else:

            def allow_keyword(keyword):
                for forbidden in {"freq", "opt", "irc", "print"}:
                    if forbidden in keyword.lower():
                        return False
                return True

            keywords = [
                keyword
                for keyword in data.metadata["keywords"]
                if allow_keyword(keyword)
            ]

            method = " ".join(keywords)

        solvent = args.solvent
        blocks = f"%pal\n nprocs {args.nprocs}\nend\n%scf\n maxiter {args.iterations}\nend"
        if solvent != "none" and not args.pm3:
            blocks += f'\n%cpcm\n smd true\n smdsolvent "{solvent}"\nend'

        if "ORCA_COMMAND" not in os.environ:
            # For parallel runs ORCA has to be called with full pathname
            os.environ["ORCA_COMMAND"] = shutil.which("orca")

        calc = ORCA(
            label="012345_swing", orcasimpleinput=method, orcablocks=blocks
        )

    print(f"*** {method} ***")
    print(f"    : solvent:              {solvent}")

    atoms.set_calculator(calc)
    potential_min = atoms.get_potential_energy()
    print(f"@ potential energy:        {potential_min} eV")

    indices = np.where(data.vibfreqs < 0)[0]
    n_indices = len(indices)
    print(f"@ imaginary frequencies:   {data.vibfreqs[indices]}")
    if not n_indices:
        print("    : nothing to be done, bye")
        return

    ignoring = None
    if args.transition_state:
        ignoring = 0
        print("    : transition state:    ignoring first imaginary frequency")

    omegas = []
    potentials = []

    def f(omega):
        atoms.set_positions(
            initial_positions
            + np.einsum("i,ijk->jk", omega, data.vibdisps[indices])
        )

        potential = 1e3 * (atoms.get_potential_energy() - potential_min)

        omegas.append(omega)
        potentials.append(potential)
        print(f"    : omega:               {omega}")
        print(f"    : potential:           {potential} meV")

        return potential

    if args.minimize:
        guesses = [np.zeros_like(indices, dtype=float)]

    for i in indices:
        if ignoring is not None and i == ignoring:
            continue

        print(f"@ searching in direction   #{i}")

        def g(w):
            z = np.zeros_like(indices, dtype=float)
            z[i] = w
            return f(z)

        if args.minimize:
            res = minimize_scalar(
                g,
                method="bounded",
                bounds=(-args.max_omega, args.max_omega),
                tol=args.tol,
            )
            print(res)

            guess = np.zeros_like(indices, dtype=float)
            guess[i] = res.x
            guesses.append(guess)
        else:
            dx = args.max_omega / 100
            x = [-dx, 0.0, dx]
            y = [g(-dx), 0.0, g(dx)]

            # p[0] * x**2 + p[1] * x + p[2] == k * (x - x0)**2 == k * x**2 - 2 * x0 * k * x + k * x0**2
            p = np.polyfit(x, y, 2)
            print(p)
            print(np.roots(p))

            dp = np.polyder(p)
            print(dp)
            r = np.roots(dp)
            print(r)

            # k = p[0]
            # x0 = np.sqrt(p[2] / k)
            # print(k, x0)
            # print(root(lambda z: [p[0] - z[0], p[1] + 2 * z[0] * z[1], p[2] - z[0] * z[1] ** 2], [k, x0]))

            best_positions = initial_positions + np.einsum(
                "i,ijk->jk", r, data.vibdisps[indices]
            )

    if args.minimize:
        print("@ choosing initial guess for global search")
        if n_indices > 1:
            guesses.append(np.sum(guesses, axis=0))
        x0 = guesses[np.argmin([f(guess) for guess in guesses])]

        print("@ searching in all directions")
        constraints = ()
        if args.transition_state and ignoring is not None:
            constraints = (
                {"type": "eq", "fun": lambda omega: omega[ignoring]},
            )
        res = minimize(
            f,
            x0=x0,
            bounds=n_indices * [(-args.max_omega, args.max_omega)],
            constraints=constraints,
            tol=args.tol,
        )
        print(res)
        best_positions = initial_positions + np.einsum(
            "i,ijk->jk", res.x, data.vibdisps[indices]
        )

        # TODO(schneiderfelipe): correct for when using --transition-state
        omegas = np.array(omegas)
        fig, ax = plt.subplots(n_indices, 1)
        if n_indices == 1:
            ax = [ax]
        xlim = (-args.max_omega - 0.05, args.max_omega + 0.05)
        ylim = (np.min(potentials) - 2.0, 40.0)
        for i in indices:
            if ignoring is not None and i == ignoring:
                continue

            ax[i].plot(omegas[:, i], potentials, "o")
            ax[i].set_title(f"view of normal mode #{i}")
            ax[i].set_ylabel(r"potential energy (meV)")
            ax[i].set_xlabel(rf"$\omega_{i}$")
            ax[i].set_ylim(ylim)
            ax[i].set_xlim(xlim)
        plt.tight_layout()
        plt.show()

    print("@ writing best geometry to swinged.xyz")
    # TODO(schneiderfelipe): print a RMSD between initial and final structures
    atoms.set_positions(best_positions)
    atoms.write("swinged.xyz", format="xyz", plain=True)

if __name__ == "__main__":
    main()


from utils import *
import matplotlib.pyplot as plt
from utils import rel_rot, rot_vec
from numpy import savetxt
import numpy as np
import scipy
from scipy.spatial import distance
from scipy.spatial.transform import Rotation as R
from scipy import stats, optimize, interpolate
import random
from Networks import NavACLNetwork, SoftQNetwork
from isaac_q_helper_functions import *
from RobotTask import RobotTask, Tasktype
import json
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque
from per_buffer import ReplayBuffer
import random
import time

class Nav_ACL():
  """
    My Interpretation of the Nav_ACL Algorithm for automatic Curriculum generation
    https://arxiv.org/abs/2009.05429
  """
  def __init__(self,nav_acl_network,config,worker_instance=True, shared_q_net=None, env=None, task_offset=np.array([0,0,0]), agent_index = 0, shared_elements=None):
    super(Nav_ACL, self).__init__()
    """
        my Implementaion keeps track of all previously generated tasks and their outcomes
    """
    self.config                         = config
    self.tasks                          = []
    self.steps                          = []
    self.rewards                        = []
    self.dones                          = []
    self.params                         = []
    self.default_task                   = self.config['default_task_omniverse']
    self.GOID_limits                    = self.config['GOID_limits']
    self.task_gen_method                = self.config['task_generation_method']
    self.initial_orientation            = np.array([1,0]) # it is assumed that the initial 2D orientation of the robot and the dolly is [1,0] and that the applied is relative to this initial orientation
    self.NavACL_loss_func               = torch.nn.BCELoss()
    self.adaptive_filtering_task_types  = [Tasktype.RANDOM, Tasktype.EASY, Tasktype.FRONTIER]
    self.adaptive_filtering_task_probs  = [0.15  ,   0.425   ,   0.425 ]
    self.device                         = train_device
    self.shared_elements                = shared_elements
    self.agent_index                    = agent_index
    self.q_net                          = SoftQNetwork(self.config)
    self.q_net.to(Qricculum_device)

    if(not worker_instance): # this instance of NavACL is the one used to train the network i.e. it does not receive copies ->  it produces the copies
      self.NavACLNet        = nav_acl_network
      self.NavACL_optimizer = optim.Adam(self.NavACLNet.parameters(), lr=self.config['nav_acl_lr'])
      self.q_net = None
      self.env   = None
    else:
      if(self.config['q_ricculum_learning']):
        if(self.config['use_AF_database']): ## use a predefined database for adaptive filtering method of nav_acl (only relevant for qricculum learning, bc. it prevents generating tasks wich are only used for creating the adaptive boundaries)
          self.load_AF_database(self.config['AF_database_path']+str(agent_index)+".npy")
          print("using qricculum learning for agent: ", agent_index, " with task AF_database: ", self.config['AF_database_path']+str(agent_index)+".npy")
        assert (shared_q_net is not None)
        assert (env is not None)
        self.env = env
        self.update_networks(nav_acl_network,shared_q_net,task_offset)
      else:
        self.update_networks(nav_acl_network,q_net=None)
        self.NavACLNet = nav_acl_network

    if(self.config['using_unity']):
      self.default_task  = self.config['default_task_unity']
    else:
      self.default_task = self.config['default_task_omniverse']

  def update_networks(self,nav_acl_net, q_net, translation_offset=np.array([0,0,0])):
    if(self.config['q_ricculum_learning']):
        assert (q_net is not None)
        self.q_net.load_state_dict(q_net.state_dict())
        for param in self.q_net.parameters():
            param.requires_grad = False

    if(self.config['using_unity']):
      self.default_task  = self.config['default_task_unity']

    self.NavACLNet = NavACLNetwork(5,self.config['nav_acl_hidden_dim'])
    self.NavACLNet.load_state_dict(nav_acl_net.state_dict())
    self.NavACLNet.to(self.device)

  def get_task_difficulty_measure(self, robot_task):
    success_probability = self.get_task_success_probability(robot_task)
    return (1-success_probability)

  def get_task_success_probability(self, robot_task):
    task_params_array = self.get_task_params_array(robot_task,normalize=self.config['normalize_tasks'])
    task_params_array = torch.FloatTensor(task_params_array).to(self.device)
    return self.NavACLNet(task_params_array)

  def normalize_task_params_array(self, task_params_array):
    # task_params_array = (task_params_array - self.task_mean) / (self.deviser+1e-3)
    r_rot = self.config['randomization_params']['robot_randomization']['rotation_rnd'][0] / 2
    d_rot = self.config['randomization_params']['dolly_randomization']['rotation_rnd'][0] / 2
    rot   = r_rot + d_rot
    task_params_array[0] = np.interp(task_params_array[0], (self.config['randomization_params']['min_dist_dolly_robot'], self.config['randomization_params']['max_dist_dolly_robot']), (-1,1))
    task_params_array[1] = np.interp(task_params_array[1], (self.config['randomization_params']['min_dist_dolly_obs'], self.config['randomization_params']['max_dist_dolly_obs']), (-1,1))
    task_params_array[2] = np.interp(task_params_array[2], (self.config['randomization_params']['min_dist_robot_obs'], self.config['randomization_params']['max_dist_robot_obs']), (-1,1))
    task_params_array[3] = np.interp(task_params_array[3], (-(rot*np.pi)/180, (rot*np.pi)/180), (-1,1))
    task_params_array[4] = np.interp(task_params_array[4], (self.config['collision_penalty']-self.config['goal_reward'],self.config['goal_reward']), (-1,1))
    return task_params_array

  def get_task_params_array(self, task,normalize=True):
    if(self.config['q_ricculum_learning']):
      task_params_array = np.array([self.compute_dolly_robot_distance(task),
                                    self.compute_dolly_min_distance(task),
                                    self.compute_robot_min_distance(task),
                                    self.compute_relative_rotation_robot_dolly(task),
                                    np.clip(task.q_value,self.config['collision_penalty']-self.config['goal_reward'],self.config['goal_reward'])])
    else:
      task_params_array = np.array([self.compute_dolly_robot_distance(task),
                                    self.compute_dolly_min_distance(task),
                                    self.compute_robot_min_distance(task),
                                    self.compute_relative_rotation_robot_dolly(task),
                                    0.1])
    if normalize:
      task_params_array = self.normalize_task_params_array(task_params_array)
    return task_params_array

  def apply_offset_to_robot(self, robot_task, translation_offset):
    robot_task.robot_translation      =robot_task.robot_translation + translation_offset
    return robot_task

  def apply_offset_to_dolly(self, robot_task, translation_offset):
    robot_task.dolly_translation      += translation_offset
    return robot_task

  def create_randomized_task(self,translation_offset= np.array([0,0,0])):
    robot_task                        = RobotTask(self.default_task)
    robot_task.robot_translation      =robot_task.robot_translation + translation_offset
    robot_task.dolly_translation      += translation_offset
    robot_task.obstacle_translation   += translation_offset
    robot_task.obstacle_1_translation += translation_offset
    robot_task.obstacle_2_translation += translation_offset
    robot_task.obstacle_3_translation += translation_offset
    robot_task.randomize_task(self.config['randomization_params'])
    return robot_task

  def sample_random_task(self,translation_offset= np.array([0,0,0])):
    robot_task = self.create_randomized_task(translation_offset=translation_offset)

    while self.check_if_task_is_valid(robot_task) is not True:
      robot_task = self.create_randomized_task(translation_offset=translation_offset)

    if(self.config['q_ricculum_learning']):
      robot_task.q_value = self.get_q_value_for_task(robot_task,translation_offset)
      # print("q value of task: ", self.get_task_params_array(robot_task,False), " is : ", robot_task.q_value)
    return robot_task

  def save_tasks_for_later(self,translation_offset, num_tasks=1000, save_location="/home/developer/Training_results/Qricculum_Learning"):
    task_list             = []
    for task in range(num_tasks):
      robot_task            = self.create_randomized_task(translation_offset=translation_offset)
      obs_cam, obs_lidar    = self.env.reset(robot_task)
      prev_action           = np.zeros((2))
      prev_reward           = [self.config['step_penalty']]
      task_array            = robot_task.get_task_array()
      task_list.append((obs_cam, obs_lidar,prev_action,prev_reward,task_array))
      if (task % 1000 == 0 ):
        print("saved so far: ", task)
    task_array = np.array(task_list)
    np.save(save_location,task_array)

  def load_AF_database(self, AF_database_path):
    """
      loads a task database from an .npy file and stores it as a list
    """
    self.AF_database = np.load(AF_database_path, allow_pickle=True)
    self.AF_database = self.AF_database.tolist() # this way we can use random.sample(list, num_elements)
    print("using qricculum learning with database: ", AF_database_path)
    self.dtb_stack = self.create_dtb_stack()

  def observation_to_q_value(self, obs_cam, obs_lidar, prev_action, prev_reward):
    n_s_f = self.config['num_stacked_frames']
    stacked_camera_obsrv  = deque(maxlen=n_s_f)
    transitions           = deque(maxlen=n_s_f)
    # start with a fresh state
    for i in range(self.config['num_stacked_frames']):
        stacked_camera_obsrv.append(obs_cam)
        if(self.config['use_snail_mode']):
            transitions.append((obs_cam, obs_lidar,prev_action,prev_reward))

    stacked_camera_state = np.asarray(stacked_camera_obsrv).reshape((self.env.observtion_shape[0]*n_s_f,self.env.observtion_shape[1],self.env.observtion_shape[2]))
    cam_tensor           = torch.FloatTensor(stacked_camera_state).unsqueeze(0).to(Qricculum_device)
    lidar_tensor         = torch.FloatTensor(obs_lidar).unsqueeze(0).to(Qricculum_device)
    prev_action          = torch.FloatTensor(prev_action).unsqueeze(0).to(Qricculum_device)
    reward               = torch.FloatTensor(prev_reward).unsqueeze(0).to(Qricculum_device)       # reward is single value, unsqueeze() to add one dim to be [reward] at the sample dim;
    with torch.no_grad():
      q_value = self.q_net(cam_tensor,lidar_tensor, prev_action, reward).detach().cpu().numpy().flatten()
      del stacked_camera_state, cam_tensor, lidar_tensor, prev_action, reward, obs_cam, obs_lidar, stacked_camera_obsrv
      return q_value[0]

  def get_q_value_for_task(self,robot_task,translation_offset):
    obs_cam, obs_lidar    = self.env.reset(robot_task)
    prev_action           = np.zeros(8)
    prev_reward           = np.ones(4)*self.config['step_penalty']
    # self.show_state(obs_cam,True,string_extra=str(translation_offset))
    q                     = self.observation_to_q_value(obs_cam, obs_lidar,prev_action, prev_reward)
    return q

  def show_state(self,state_tensor, save=False, string_extra=""):
      state_shape= (84,84,3)
      state_tensor = state_tensor / 2**8
      pic0 = np.array(state_tensor.reshape(state_shape).astype(np.float32))

      fig, (ax1) = plt.subplots(1, 1)
      ax1.imshow(pic0)
      if(save):
        path = "/home/developer/Pictures"+string_extra+str(current_milli_time())+".png"
        print("saving image as: ", path)
        plt.imsave(path,pic0)
      else:
        plt.show()

  def generate_random_task(self, translation_offset=np.array([0,0,0])):
    """
      returns a randomized RobotTask:

      based on a default task given in NAV_ACL_hyperparameters.json
    """
    if self.task_gen_method == "GOID":
      return self.generate_random_task_GOID(translation_offset)

    if self.task_gen_method == "AF":
      return self.generate_random_task_AF(translation_offset)

  def generate_random_task_GOID(self, translation_offset):
    while True:
      robot_task = self.sample_random_task(translation_offset)
      difficulty = self.get_task_difficulty_measure(robot_task).cpu()
      if ( (difficulty > self.GOID_limits['lower_limit'] ) and (difficulty < self.GOID_limits['upper_limit'])) :
        return robot_task

  def generate_random_task_AF(self, translation_offset):
    if(self.config['only_random_tasks']):
      task_type = Tasktype.RANDOM
    else:
      task_type = np.random.choice(self.adaptive_filtering_task_types, 1, p=self.adaptive_filtering_task_probs)[0] # adaptive sampling of the task type
    return self.get_dynamic_task(task_type,translation_offset)

  def get_dynamic_task(self, task_type, translation_offset):
    """
      This is the main algorithm for nav-acl-q
    """

    if(self.config['use_AF_database']):
      mu, sigma = self.fast_create_adaptive_boundaries_mu_sig()
    else:
      mu, sigma = self.create_adaptive_boundaries()

    beta      = self.config['adaptive_filtering_params']['nav_beta']
    gamma_low = self.config['adaptive_filtering_params']['nav_gamma_low']
    gamma_hi  = self.config['adaptive_filtering_params']['nav_gamma_hi']
    omega     = self.config['adaptive_filtering_params']['nav_P_omega']

    for trial in range(self.config['nav_acl_max_AF_task_samples']):
      if(self.config['runtimetasks_from_database'] and self.config['use_AF_database']):
        robot_task = self.sample_task_from_database(self.AF_database)
      else:
        robot_task = self.sample_random_task(translation_offset)
      succ_prob  = self.get_task_success_probability(robot_task).cpu()
      time.sleep(0.01)

      if task_type == Tasktype.EASY:
        if (mu + beta*sigma) < 1:
          if (succ_prob > (mu + beta*sigma)) or (succ_prob > omega):
            robot_task.task_type = Tasktype.EASY
            return robot_task
        else:
          if (succ_prob > mu):
            robot_task.task_type = Tasktype.EASY
            return robot_task
      if task_type == Tasktype.FRONTIER:
        if (mu - gamma_low*sigma) < succ_prob < (mu + gamma_hi*sigma):
          robot_task.task_type = Tasktype.FRONTIER
          return robot_task
      if task_type == Tasktype.RANDOM:
        robot_task.task_type = Tasktype.RANDOM
        return robot_task
    print("TOOK too many trials to find task of type: ", task_type, " with boundaries mu: ", mu, " sigma: ", sigma)
    robot_task.task_type = Tasktype.RANDOM
    return robot_task

  def create_adaptive_boundaries(self, translation_offset=np.array([0,0,0])):
    tasks_parameters = []
    if(self.config['use_AF_database']):
      tasks_parameters = self.get_task_parameters_for_AF_database()

    else:
      for i in range(self.config['nav_acl_AF_boundaries_task_samples']):
        tasks_parameters.append(self.get_task_params_array(self.sample_random_task(translation_offset=translation_offset),self.config['normalize_tasks']))

    params      = torch.FloatTensor(tasks_parameters).to(self.device)
    predictions = self.NavACLNet(params).cpu().detach().numpy()
    mu, sigma   = scipy.stats.distributions.norm.fit(predictions)
    return mu, sigma

  def get_task_parameters_for_AF_database(self):
    tasks_parameters = []
    for i in range(self.config['nav_acl_AF_boundaries_task_samples']):
      task = self.sample_task_from_database(self.AF_database)
      tasks_parameters.append(self.get_task_params_array(task,self.config['normalize_tasks'])) # and append it to the task array which will be used for creating the adaptive boundaries
    return tasks_parameters

  def sample_task_from_database(self,database):
    task_sample                        = random.sample(database, 1)[0]
    obs_cam, obs_lidar,_,_,task_array  = task_sample
    prev_action                        = np.zeros(8)
    prev_reward                        = np.ones(4)*self.config['step_penalty']
    q                                  = self.observation_to_q_value(obs_cam,obs_lidar,prev_action,prev_reward) # pass observation through the q network to obtain currently estimated q value
    task                               = RobotTask(self.default_task)                                           # create a task instance
    task.from_task_array(task_array)                                                                            # initialize according to the task that was sampled from the database
    task.q_value = q                                                                                # replace the q
    return task

  def compute_dolly_robot_distance(self, task):
      return distance.euclidean(task.robot_translation, task.dolly_translation)

  def compute_obstacle_distance_dolly(self, task):
      distances = np.array([distance.euclidean(task.obstacle_2_translation, task.dolly_translation),
                            distance.euclidean(task.obstacle_3_translation, task.dolly_translation)])
      return np.min(distances)

  def compute_obstacle_distance_robot(self, task):
      distances = np.array([distance.euclidean(task.obstacle_translation,   task.robot_translation),
                            distance.euclidean(task.obstacle_1_translation, task.robot_translation)])
      return np.min(distances)

  def compute_path_complexity(self, task):
      return 0.1

  def check_if_task_is_valid(self,task):
    # first check is no object is too close to the dolly:
    if(self.compute_dolly_min_distance(task) < self.config['randomization_params']['min_dist_dolly_obs']):
      return False
    elif(self.compute_robot_min_distance(task) < self.config['randomization_params']['min_dist_robot_obs']):
      return False
    else:
      return True

  def compute_robot_min_distance(self,task):
    return self.compute_min_distance(task, task.robot_translation)
  def compute_dolly_min_distance(self,task):
    return self.compute_min_distance(task, task.dolly_translation)

  def compute_min_distance(self,task, target_translation):
    distances = np.zeros((4))
    obstacles = task.get_obstacle_translations_array()
    for i in range(4):
      distances[i] = distance.euclidean(obstacles[i], target_translation)
    min_dist = np.min(distances)

    if min_dist > 50:
      return 0
    else :
      return np.min(distances)

  def compute_relative_rotation_robot_dolly(self, task, as_degrees=False):
    Robot_orientation     = rot_vec(self.initial_orientation,task.robot_rotation[0])
    Dolly_orientation     = rot_vec(self.initial_orientation,task.dolly_rotation[0])
    relative_rotation_rad = rel_rot(Robot_orientation,Dolly_orientation)
    if(as_degrees):
      return relative_rotation_rad * (180/np.pi)
    else:
      return relative_rotation_rad

  def train(self,robot_task,label,batch_mode = False):
    self.NavACL_optimizer.zero_grad()

    if(self.config['nav_acl_batch_mode']):
      params   = []
      labels   = []
      for t in range(self.config['nav_acl_batch_size']):
        if(type(robot_task[t])== RobotTask):
          params.append(self.get_task_params_array(robot_task[t],self.config['normalize_tasks']))
        else:
          params.append(self.normalize_task_params_array(robot_task[t]))
        labels.append(label[t])

      task_params_array = torch.FloatTensor(params).to(self.device)
      label             = torch.FloatTensor(labels).flatten().to(self.device)

    else:
      if(type(robot_task)== RobotTask):
        task_params_array = self.get_task_params_array(robot_task,self.config['normalize_tasks'])
      else:
        task_params_array = self.normalize_task_params_array(robot_task)

      task_params_array = torch.FloatTensor(task_params_array).flatten().to(self.device)
      label             = torch.FloatTensor(label).flatten().to(self.device)

    prediction = self.NavACLNet(task_params_array)
    loss = self.NavACL_loss_func(prediction, label)

    loss.backward()
    self.NavACL_optimizer.step()
    return prediction.detach().cpu().numpy()

  def batch_train(self,robot_tasks, labels):
    self.NavACL_optimizer.zero_grad()

    params   = []
    _labels   = []
    for t in range(self.config['nav_acl_batch_size']):
      if(type(robot_tasks[t])== RobotTask):
        params.append(self.get_task_params_array(robot_tasks[t],self.config['normalize_tasks']))
      else:
        params.append(robot_tasks[t])

      _labels.append(labels[t])

    task_params_array = torch.FloatTensor(params).to(self.device)
    label             = torch.FloatTensor(_labels).flatten().to(self.device)
    prediction        = self.NavACLNet(task_params_array)
    # print("robot_tasks:" , robot_tasks)
    # print("predictions: ", prediction)
    loss              = self.NavACL_loss_func(prediction, label.view((prediction.shape))) # labels are [batchsize] abd predictions are [batchsize,1] thus labels have to be viewed as [8,1]! This is very important, otherwise it will only use one of the parameters for esitmation

    loss.backward()
    self.NavACL_optimizer.step()
    return prediction.detach().cpu().numpy()

  def create_dtb_stack(self):
      af_database = np.load(self.config['AF_database_path']+str(self.agent_index)+".npy", allow_pickle=True)
      af_database = np.array(af_database).reshape(af_database.shape[0],af_database.shape[1])
      dtb_stack = []
      for t_i in range(af_database.shape[0]):
          image_stack = np.array([af_database[t_i][0],af_database[t_i][0],af_database[t_i][0],af_database[t_i][0]]).reshape(12,80,80)
          lidar       = af_database[t_i][1]
          action      = np.zeros((8))
          reward      = np.ones(4)*self.config['step_penalty']
          task_array  = af_database[t_i][4]
          robot_task  = RobotTask(self.config['default_task_unity'])
          robot_task.from_task_array(task_array)
          task_params_array = self.get_task_params_array(robot_task,self.config['normalize_tasks'])
          task_tuple = (image_stack, lidar, action, reward, task_params_array)
          dtb_stack.append(task_tuple)
      return dtb_stack

  def fast_create_adaptive_boundaries_mu_sig(self):
    c, l, a, r, tsk = map(np.stack, zip(*random.sample(self.dtb_stack, self.config['nav_acl_AF_boundaries_task_samples'])))

    cam_tensor          = torch.FloatTensor(c).to(Qricculum_device)
    lidar_tensor        = torch.FloatTensor(l).to(Qricculum_device)
    prev_action         = torch.FloatTensor(a).to(Qricculum_device)
    prev_reward         = torch.FloatTensor(r).to(Qricculum_device)  # prev_reward is single value, unsqueeze() to add one dim to be [reward] at the sample dim;
    with torch.no_grad():
        predicted_q_values = self.q_net(cam_tensor,lidar_tensor, prev_action, prev_reward).detach().cpu().numpy()
        predicted_q_values = np.interp(predicted_q_values, (self.config['collision_penalty']-self.config['goal_reward'],self.config['goal_reward']), (-1,1))

    tsk[:,4] = predicted_q_values[:,0]
    # print(tsk)
    tsk_tensor = torch.FloatTensor(tsk).to(self.device)
    predictions = self.NavACLNet(tsk_tensor).detach().cpu().numpy()
    mu, sigma   = scipy.stats.distributions.norm.fit(predictions)
    return mu, sigma

import numpy,from numpy,import scipy,from scipy


import os
import math
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt

def py_cpu_nms(dets, scores, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    # scores = dets[:, 4]  # bbox打分

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    # 打分从大到小排列，取index
    order = scores.argsort()[::-1]
    # keep为最后保留的边框
    keep = []
    while order.size > 0:
        # order[0]是当前分数最大的窗口，肯定保留
        i = order[0]
        keep.append(i)
        # 计算窗口i与其他所有窗口的交叠部分的面积
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        # 交/并得到iou值
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        # inds为所有与窗口i的iou值小于threshold值的窗口的index，其他窗口此次都被窗口i吸收
        inds = np.where(ovr <= thresh)[0]
        # order里面只保留与窗口i交叠面积小于threshold的那些窗口，由于ovr长度比order长度少1(不包含i)，所以inds+1对应到保留的窗口
        order = order[inds + 1]

    return keep

def vis(im_path, im_x_crop, score, cls, ctr, box):
    """
    :param im_path:
    :type im_path:
    :param im_x_crop:
    :type im_x_crop:  shape=(303, 303, 3)
    :param score: 总分
    :type score: shape=(361,)
    :param cls: 分类得分
    :type cls: shape=(361, 1)
    :param ctr:
    :type ctr: shape=(361, 1)
    :param box: xyxy
    :type box: shape=(361, 4)
    :return:
    :rtype:
    """
    def reshape_resize(array):
        resize = False
        array = array.reshape(score_map_size, score_map_size) * 255
        if resize:
            array = cv2.resize(array, (image_size, image_size))
        array = array.astype(np.uint8)
        return array

    im_x_crop = cv2.cvtColor(im_x_crop, cv2.COLOR_RGB2BGR)
    video_name, im_name = im_path.split('/')[-2:]
    image_size = im_x_crop.shape[0]
    score_map_size = int(math.sqrt(len(score)))
    score = reshape_resize(score)
    cls = reshape_resize(cls[:, 0])
    ctr = reshape_resize(ctr[:, 0])

    """
    '''可视化候选框'''
    ids = py_cpu_nms(box, score, 0.3)  # ids = score.argsort()[-100:][::-1]
    top_box = box[ids]
    for b in top_box:
        x1, y1, x2, y2 = b
        im_x_crop = cv2.rectangle(im_x_crop, (x1, y1), (x2, y2),
                                  (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
                                  thickness=2)
    """

    plt.subplot(221)
    plt.imshow(im_x_crop)
    plt.subplot(222)
    plt.imshow(score)
    plt.subplot(223)
    plt.imshow(cls)
    plt.subplot(224)
    plt.imshow(ctr)
    save_dir = os.path.join('/tmp/vis', video_name)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    save_path = os.path.join(save_dir, im_name)
    plt.savefig(save_path)
    save_path_1 = os.path.join(save_dir, 'im_' + im_name)
    cv2.imwrite(save_path_1, cv2.cvtColor(im_x_crop, cv2.COLOR_RGB2BGR))
    return


"""Test cases for evaluation scripts."""
import os
import unittest

import numpy as np

from ..common.typing import NDArrayF64
from ..label.io import load, load_label_config
from ..unittest.util import get_test_file
from .tagging import evaluate_tagging

class TestScalabelTaggingEval(unittest.TestCase):
    """Test cases for Scalabel tagging evaluation."""

    cur_dir = os.path.dirname(os.path.abspath(__file__))
    gts_path = f"{cur_dir}/testcases/tagging/tag_gts.json"
    preds_path = f"{cur_dir}/testcases/tagging/tag_preds.json"
    gts = load(gts_path).frames
    preds = load(preds_path).frames
    config = load_label_config(get_test_file("tagging/tag_configs.toml"))
    result = evaluate_tagging(gts, preds, config, nproc=1)

    def test_frame(self) -> None:
        """Test case for the function frame()."""
        data_frame = self.result.pd_frame()
        categories = set(
            [
                "weather.rainy",
                "weather.snowy",
                "weather.clear",
                "weather.overcast",
                "weather.undefined",
                "weather.partly cloudy",
                "weather.foggy",
                "scene.city street",
                "scene.gas stations",
                "scene.parking lot",
                "scene.residential",
                "scene.highway",
                "scene.tunnel",
                "scene.undefined",
                "WEATHER",
                "SCENE",
                "AVERAGE",
            ]
        )
        self.assertSetEqual(categories, set(data_frame.index.values))

        data_arr = data_frame.to_numpy()
        aps: NDArrayF64 = np.array(
            [100.0, -1.0, 100.0, 100.0, 100.0, 0.0, 0.0, -1.0, -1.0, -1.0]
            + [-1.0, 100.0, -1.0, 100.0, 66.66667, 100.0, 83.33333],
            dtype=np.float64,
        )
        self.assertTrue(
            np.isclose(
                np.nan_to_num(data_arr[:, 0], nan=-1.0), aps, atol=1e-2
            ).all()
        )

        overall_scores: NDArrayF64 = np.array(
            [83.33333333, 80.55555556, 81.66666667, 90.0], dtype=np.float64
        )
        self.assertTrue(
            np.isclose(
                np.nan_to_num(data_arr[-1], nan=-1.0),
                overall_scores,
                atol=1e-2,
            ).all()
        )

    def test_summary(self) -> None:
        """Check evaluation scores' correctness."""
        summary = self.result.summary()
        overall_reference = {
            "precision": 83.33333333333333,
            "recall": 80.55555555555556,
            "f1_score": 81.66666666666666,
            "accuracy": 90.0,
            "recall/WEATHER": 61.11111111111111,
            "f1_score/SCENE": 100.0,
            "accuracy/SCENE": 100.0,
            "recall/SCENE": 100.0,
            "f1_score/WEATHER": 63.33333333333333,
            "precision/WEATHER": 66.66666666666666,
            "precision/SCENE": 100.0,
            "accuracy/WEATHER": 80.0,
        }
        self.assertSetEqual(set(summary.keys()), set(overall_reference.keys()))
        for name, score in summary.items():
            if np.isnan(score):
                score = np.nan_to_num(score, nan=-1.0)
            self.assertAlmostEqual(score, overall_reference[name])


# Copyright (C) 2018 Atsushi Togo
# All rights reserved.
#
# This file is part of phonopy.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import sys
import numpy as np

def get_fc2(supercell,
            primitive,
            displacements,
            forces,
            atom_list=None,
            options=None,
            log_level=0):
    p2s_map = primitive.p2s_map
    is_compact_fc = (atom_list is not None and
                     (atom_list == p2s_map).all())
    fc2 = run_alm(supercell,
                  primitive,
                  displacements,
                  forces,
                  1,
                  is_compact_fc=is_compact_fc,
                  options=options,
                  log_level=log_level)[0]

    if not is_compact_fc and atom_list is not None:
        fc2 = np.array(fc2[atom_list], dtype='double', order='C')

    return fc2

def run_alm(supercell,
            primitive,
            displacements,
            forces,
            maxorder,
            is_compact_fc,
            options=None,
            log_level=0):
    fcs = None  # This is returned.

    lattice = supercell.cell
    positions = supercell.scaled_positions
    numbers = supercell.numbers
    natom = len(numbers)
    p2s_map = primitive.p2s_map
    p2p_map = primitive.p2p_map

    alm_options = _update_options(options)
    num_elems = len(np.unique(numbers))
    if log_level:
        print("---------------------------------"
              " ALM start "
              "--------------------------------")
        print("ALM is a non-trivial force constants calculator. "
              "Please cite the paper:")
        print("T. Tadano and S. Tsuneyuki, "
              "J. Phys. Soc. Jpn. 87, 041015 (2018).")
        print("ALM is developed at https://github.com/ttadano/ALM by T. "
              "Tadano.")
    if log_level == 1:
        print("Use -v option to watch detailed ALM log.")

    shape = (maxorder, num_elems, num_elems)
    cutoff_radii = -np.ones(shape, dtype='double')
    if alm_options['cutoff'] is not None:
        if len(alm_options['cutoff']) == 1:
            cutoff_radii[:] = alm_options['cutoff'][0]
        elif np.prod(shape) == len(alm_options['cutoff']):
            cutoff_radii[:] = np.reshape(alm_options['cutoff'], shape)
        else:
            raise RuntimeError("Cutoff is not properly set.")

    _disps, _forces, df_msg = _slice_displacements_and_forces(
        displacements,
        forces,
        alm_options['ndata'],
        alm_options['nstart'],
        alm_options['nend'])

    if log_level > 0:
        log_level_alm = log_level - 1
    else:
        log_level_alm = 0

    try:
        from alm import ALM, optimizer_control_data_types
    except ImportError:
        raise ImportError("ALM python module was not found.")

    with ALM(lattice, positions, numbers) as alm:
        if log_level > 0:
            if alm_options['cutoff'] is not None:
                for i in range(maxorder):
                    if maxorder > 1:
                        print("fc%d" % (i + 2))
                    print(("cutoff" + " %6s" * num_elems)
                          % tuple(alm.kind_names.values()))
                    for r, kn in zip(cutoff_radii[i], alm.kind_names.values()):
                        print(("   %-3s" + " %6.2f" * num_elems)
                              % ((kn, ) + tuple(r)))
            if df_msg is not None:
                print(df_msg)
        if log_level > 1:
            print("")
        sys.stdout.flush()

        alm.output_filename_prefix = alm_options['output_filename_prefix']
        alm.verbosity = log_level_alm

        alm.define(
            maxorder,
            cutoff_radii=cutoff_radii,
            symmetrization_basis=alm_options['symmetrization_basis'])
        alm.displacements = _disps
        alm.forces = _forces

        # Mainly for elastic net (or lasso) regression
        optcontrol = {}
        for key in optimizer_control_data_types:
            if key in alm_options:
                optcontrol[key] = alm_options[key]
        if optcontrol:
            alm.optimizer_control = optcontrol
            if ('cross_validation' in optcontrol and
                optcontrol['cross_validation'] > 0):
                alm.optimize(solver=alm_options['solver'])
                optcontrol['cross_validation'] = 0
                optcontrol['l1_alpha'] = alm.cv_l1_alpha
                alm.optimizer_control = optcontrol

        alm.optimize(solver=alm_options['solver'])

        fcs = _extract_fc_from_alm(alm,
                                   natom,
                                   maxorder,
                                   is_compact_fc,
                                   p2s_map=p2s_map,
                                   p2p_map=p2p_map)

    if log_level:
        print("----------------------------------"
              " ALM end "
              "---------------------------------")

    return fcs

def _update_options(fc_calculator_options):
    """Set ALM options with appropriate data types

    fc_calculator_options : str
        This string should be written such as follows:

            "solver = dense, cutoff = 5"

        This string is parsed as collection of settings that are separated by
        comma ','. Each setting has the format of 'option = value'. The value
        is cast to have its appropriate data type for ALM in this method.

    """

    try:
        from alm import optimizer_control_data_types
    except ImportError:
        raise ImportError("ALM python module was not found.")

    # Default settings.
    alm_options = {'solver': 'SimplicialLDLT',
                   'ndata': None,
                   'nstart': None,
                   'nend': None,
                   'cutoff': None,
                   'symmetrization_basis': 'Lattice',
                   'output_filename_prefix': None}

    if fc_calculator_options is not None:
        alm_option_types = {'cutoff': np.double,
                            'ndata': int,
                            'nstart': int,
                            'nend': int,
                            'output_filename_prefix': str,
                            'solver': str,
                            'symmetrization_basis': str}
        alm_option_types.update(optimizer_control_data_types)
        for option_str in fc_calculator_options.split(","):
            key, val = [x.strip() for x in option_str.split('=')[:2]]
            if key.lower() in alm_option_types:
                if alm_option_types[key.lower()] is np.double:
                    option_value = np.array(
                        [float(x) for x in val.split()], dtype='double')
                elif alm_option_types[key.lower()] is np.intc:
                    option_value = np.array(
                        [int(x) for x in val.split()], dtype='intc')
                else:
                    option_value = alm_option_types[key.lower()](val)
                alm_options[key] = option_value
    return alm_options

def _slice_displacements_and_forces(d, f, ndata, nstart, nend):
    msg = None
    if ndata is not None:
        _d = d[:ndata]
        _f = f[:ndata]
        msg = "Number of displacement supercells: %d" % ndata
    elif nstart is not None and nend is not None:
        _d = d[nstart - 1:nend]
        _f = f[nstart - 1:nend]
        msg = "Supercell index range: %d - %d" % (nstart, nend)
    else:
        return d, f, None

    return (np.array(_d, dtype='double', order='C'),
            np.array(_f, dtype='double', order='C'), msg)

def _extract_fc_from_alm(alm,
                         natom,
                         maxorder,
                         is_compact_fc,
                         p2s_map=None,
                         p2p_map=None):
    # Harmonic: order=1, 3rd: order=2, ...
    fcs = []
    for order in range(1, maxorder + 1):
        fc = None
        p2s_map_alm = alm.getmap_primitive_to_supercell()[0]
        if (is_compact_fc and
            len(p2s_map_alm) == len(p2s_map) and
            (p2s_map_alm == p2s_map).all()):
            fc_shape = (
                (len(p2s_map), ) + (natom, ) * order + (3, ) * (order + 1))
            fc = np.zeros(fc_shape, dtype='double', order='C')
            for fc_elem, indices in zip(
                    *alm.get_fc(order, mode='origin', permutation=1)):
                v = indices // 3
                c = indices % 3
                selection = (p2p_map[v[0]], ) + tuple(v[1:]) + tuple(c)
                fc[selection] = fc_elem

        if fc is None:
            if is_compact_fc:
                atom_list = p2s_map
            else:
                atom_list = np.arange(natom, dtype=int)
            fc_shape = (
                (len(atom_list), ) + (natom, ) * order + (3, ) * (order + 1))
            fc = np.zeros(fc_shape, dtype='double', order='C')
            for fc_elem, indices in zip(
                    *alm.get_fc(order, mode='all', permutation=1)):
                v = indices // 3
                idx = np.where(atom_list == v[0])[0]
                if len(idx) > 0:
                    c = indices % 3
                    selection = (idx[0], ) + tuple(v[1:]) + tuple(c)
                    fc[selection] = fc_elem

        fcs.append(fc)

    return fcs