import os
import re
import logging
logger = logging.getLogger(__name__)

import numpy as np
import scipy.interpolate as intp
from scipy.signal import savgol_filter
import astropy.io.fits as fits
import matplotlib.pyplot as plt

from ...echelle.imageproc import combine_images
from ...echelle.trace import find_apertures, load_aperture_set
from ...echelle.background import simple_debackground
from ...echelle.extract import extract_aperset
from ...echelle.flat import get_slit_flat
from ...utils.obslog import parse_num_seq
from ...utils.onedarray import iterative_savgol_filter
from .common import (print_wrapper, parse_3ccd_images, mosaic_3_images,
                    BiasFigure, TraceFigure)

def reduce_post2004(config, logtable):
    """
    """

    # extract keywords from config file
    section = config['data']
    rawpath     = section.get('rawpath')

    section = config['reduce']
    midpath     = section.get('midpath')
    odspath     = section.get('odspath')
    figpath     = section.get('figpath')
    mode        = section.get('mode')
    fig_format  = section.get('fig_format')
    oned_suffix = section.get('oned_suffix')
    ncores      = section.get('ncores')

    # create folders if not exist
    if not os.path.exists(figpath): os.mkdir(figpath)
    if not os.path.exists(odspath): os.mkdir(odspath)
    if not os.path.exists(midpath): os.mkdir(midpath)

    # determine number of cores to be used
    if ncores == 'max':
        ncores = os.cpu_count()
    else:
        ncores = min(os.cpu_count(), int(ncores))

    nccd = 3

    ########################## load file selection #############################
    sel_lst = {}
    filesel_filename = 'file_selection.txt'
    if os.path.exists(filesel_filename):
        sel_file = open(filesel_filename)
        for row in sel_file:
            row = row.strip()
            if len(row)==0 or row[0] in '#':
                continue
            g = row.split(':')
            key, value = g[0].strip(), g[1].strip()
            if len(value)>0:
                sel_lst[key] = value
        sel_file.close()

    ################################ parse bias ################################

    #bias_lst, bias_card_lst = get_bias_post2004(config, logtable)

    section = config['reduce.bias']
    bias_file = section['bias_file']

    if mode=='debug' and os.path.exists(bias_file):
        # load bias data from existing file
        hdu_lst = fits.open(bias_file)
        # pack bias image
        bias_lst = [hdu_lst[iccd+1+nccd].data for iccd in range(nccd)]
        hdu_lst.close()

        reobj = re.compile('GAMSE BIAS[\s\S]*')
        # filter header cards that match the above pattern
        bias_card_lst = [(card.keyword, card.value) for card in head.cards
                            if reobj.match(card.keyword)]

        message = 'Load bias data from file: "{}"'.format(bias_file)
        logger.info(message)
        print(message)

    else:

        bias_data_lst = [[] for iccd in range(nccd)]
        bias_card_lst = []
        bias_items = list(filter(lambda item: item['object'].lower()=='bias',
                                 logtable))
        n_bias = len(bias_items)

        fmt_str = '  - {:>7s} {:^17} {:^20s} {:^7}'
        head_str = fmt_str.format('frameid', 'FileID', 'Object', 'exptime')

        print('* Combine Bias Images: {}'.format(bias_file))
        print(head_str)

        for logitem in bias_items:

            fname = '{}.fits'.format(logitem['fileid'])
            filename = os.path.join(rawpath, fname)
            hdu_lst = fits.open(filename)
            data_lst, mask_lst = parse_3ccd_images(hdu_lst)
            hdu_lst.close()

            for iccd in range(nccd):
                bias_data_lst[iccd].append(data_lst[iccd])

            # print info
            message = fmt_str.format(
                        '[{:d}]'.format(logitem['frameid']),
                        logitem['fileid'], logitem['object'],
                        logitem['exptime']
                        )
            print(message)

        prefix = 'HIERARCH GAMSE BIAS '
        bias_card_lst.append((prefix + 'NFILE', n_bias))

        combine_mode = 'mean'
        section = config['reduce.bias']
        cosmic_clip  = section.getfloat('cosmic_clip')
        maxiter      = section.getint('maxiter')
        maskmode    = (None, 'max')[n_bias>=3]

        bias_card_lst.append((prefix+'COMBINE_MODE', combine_mode))
        bias_card_lst.append((prefix+'COSMIC_CLIP',  cosmic_clip))
        bias_card_lst.append((prefix+'MAXITER',      maxiter))
        bias_card_lst.append((prefix+'MASK_MODE',    str(maskmode)))

        # create the hdu list to be saved
        hdu_lst = fits.HDUList()
        # create new FITS Header for bias
        head = fits.Header()
        for card in bias_card_lst:
            head.append(card)
        head['HIERARCH GAMSE FILECONTENT 0'] = 'NONE'

        # pack the primary HDU
        hdu_lst.append(fits.PrimaryHDU(header=head))

        # prepare the list for image data of each CCD
        bias_lst = []

        # scan for each ccd
        for iccd in range(nccd):
            ### 3 CCDs loop begins here ###
            bias_data_lst[iccd] = np.array(bias_data_lst[iccd])

            sub_bias = combine_images(bias_data_lst[iccd],
                mode        = combine_mode,
                upper_clip  = cosmic_clip,
                maxiter     = maxiter,
                maskmode    = maskmode,
                ncores      = ncores,
                )

            # pack bias of each CCD into sub_bias_lst
            bias_lst.append(sub_bias)

            message = ('\033[{}mCombined bias for CCD {}: '
                        'Mean = {:6.2f}\033[0m'.format(
                        (34, 32, 31)[iccd], iccd+1, sub_bias.mean()))
            print(message)

            key = 'HIERARCH GAMSE FILECONTENT {}'.format(iccd+1)
            hdu_lst[0].header[key] = 'BIAS COMBINED FOR CCD{}'.format(iccd+1)
            hdu_lst.append(fits.ImageHDU(data=sub_bias))

        # initialize bias figure
        bias_fig = BiasFigure(data_lst=bias_lst)

        new_bias_lst = []

        # calculate and plot ymean of bias data
        for iccd in range(nccd):
            data = bias_lst[iccd]
            ny, nx = data.shape
            ally = np.arange(ny)
            ymean = data.mean(axis=1)
            ax_ycut = bias_fig.ax_ycut_lst[iccd]
            ax_ycut.plot(ymean, ally, color='C3', lw=0.5, alpha=0.6)

            # use iterative savgol filter to smooth ymean
            ysmo, _, _, _ = iterative_savgol_filter(ymean, winlen=301, order=3,
                                maxiter=10, upper_clip=3, lower_clip=3)

            ax_ycut.plot(ysmo, ally, color='r', lw=0.5, alpha=1)

            new_bias = np.tile(ysmo, (nx,1)).T
            new_bias_lst.append(new_bias)

        # save and close the figure
        figfilename = os.path.join(figpath, 'bias.png')
        bias_fig.savefig(figfilename)
        plt.close(bias_fig)

        for iccd in range(nccd):
            hdu_lst.append(fits.ImageHDU(data=new_bias_lst[iccd]))
            # update the file content in primary HDU
            key = 'HIERARCH GAMSE FILECONTENT {}'.format(1+nccd+iccd)
            card = (key, 'BIAS USED FOR CCD{}'.format(iccd+1))
            hdu_lst[0].header.append(card)
        bias_lst = new_bias_lst

        # save to FITS
        hdu_lst.writeto(bias_file, overwrite=True)

        message = 'Bias image written to "{}"'.format(bias_file)
        logger.info(message)
        print(message)

    ########################## find flat groups #########################
    flat_file = config['reduce.flat'].get('flat_file')

    flatdata_lst = []
    # a list of 3 combined flat images. [Image1, Image2, Image3]
    # bias has been corrected already. but not rotated yet.
    flatmask_lst = []
    # a list of 3 flat masks

    if mode=='debug' and os.path.exists(flat_file):
        # read flat data from existing file
        hdu_lst = fits.open(flat_file)
        for iccd in range(nccd):
            flatdata_lst.append(hdu_lst[iccd*2+1].data)
            flatmask_lst.append(hdu_lst[iccd*2+2].data)
        flatdata = hdu_lst[nccd*2+1].data.T
        flatmask = hdu_lst[nccd*2+2].data.T
        hdu_lst.close()
        message = 'Loaded flat data from file: {}'.format(flat_file)
        print(message)

        # alias of flat data and mask
        flatdata1 = flatdata_lst[0].T
        flatmask1 = flatmask_lst[0].T
        flatdata2 = flatdata_lst[1].T
        flatmask2 = flatmask_lst[1].T
        flatdata3 = flatdata_lst[2].T
        flatmask3 = flatmask_lst[2].T

    else:
        fmt_str = ('  - {:>5s} {:^17} {:<20s} {:^1s}I2 {:>7}'
                    ' \033[34m{:>8}\033[0m'
                    ' \033[32m{:>8}\033[0m'
                    ' \033[31m{:>8}\033[0m'
                    ' \033[34m{:>8}\033[0m'
                    ' \033[32m{:>8}\033[0m'
                    ' \033[31m{:>8}\033[0m')
        head_str = fmt_str.format('FID', 'fileid', 'object', '', 'exptime',
                    'nsat_1','nsat_2', 'nsat_3', 'q95_1', 'q95_2', 'q95_3')
        print('*'*10 + 'Parsing Flat Fieldings' + '*'*10)
        # print the flat list
        print(head_str)
        for logitem in logtable:
            if len(logitem['object'])>=8 and logitem['object'][0:8]=='flatlamp':
                message = fmt_str.format(
                        '[{:d}]'.format(logitem['frameid']),
                        logitem['fileid'], logitem['object'], logitem['i2'],
                        logitem['exptime'],
                        logitem['nsat_1'], logitem['nsat_2'], logitem['nsat_3'],
                        logitem['q95_1'], logitem['q95_2'], logitem['q95_3'],
                        )
                print(print_wrapper(message, logitem))

        flat_group_lst = {}
        for iccd in range(nccd):

            key = 'flat CCD%d'%(iccd+1)
            sel_string = sel_lst[key] if key in sel_lst else ''
            prompt = '\033[{1}mSelect flats for CCD {0} [{2}]: \033[0m'.format(
                      iccd+1, (34, 32, 31)[iccd], sel_string)

            # read selected files from terminal
            while(True):
                input_string = input(prompt)
                if len(input_string.strip())==0:
                    # nothing input
                    if key in sel_lst:
                        # nothing input but already in selection list
                        flat_group_lst[iccd] = parse_num_seq(sel_lst[key])
                        break
                    else:
                        # repeat prompt
                        continue
                else:
                    # something input
                    frameid_lst = parse_num_seq(input_string)
                    # pack
                    flat_group_lst[iccd] = frameid_lst
                    # put input string into selection list
                    sel_lst[key] = input_string.strip()
                    break

        # now combine flat images

        flat_hdu_lst = [fits.PrimaryHDU()]
        # flat_hdu_lst is the final HDU list to be saved as fits

        for iccd in range(nccd):
            frameid_lst = flat_group_lst[iccd]

            # now combine flats for this CCD
            flat_data_lst = []
            # flat_data_lst is a list of flat images to be combined.
            # flat_data_lst = [Image1, Image2, Image3, Image4, ... ...]

            #scan the logtable
            # log loop inside the CCD loop because flats for different CCDs are
            # in different files
            for logitem in logtable:
                if logitem['frameid'] in frameid_lst:
                    filename = os.path.join(rawpath, logitem['fileid']+'.fits')
                    hdu_lst = fits.open(filename)
                    data_lst, mask_lst = parse_3ccd_images(hdu_lst)
                    hdu_lst.close()

                    # correct bias and pack into flat_data_lst
                    if bias_lst is None:
                        flat_data_lst.append(data_lst[iccd])
                        message = 'No bias. skipped bias correction'
                    else:
                        flat_data_lst.append(data_lst[iccd]-bias_lst[iccd])
                        message = 'Bias corrected'
                    logger.info(message)

                    # print info
                    # initialize flat mask
                    if len(flat_data_lst) == 1:
                        flatmask = mask_lst[iccd]
                    flatmask = flatmask | mask_lst[iccd]

            n_flat = len(flat_data_lst)

            if n_flat == 0:
                continue
            elif n_flat == 1:
                flatdata = flat_data_lst[0]
            else:
                flat_data_lst = np.array(flat_data_lst)
                flatdata = combine_images(flat_data_lst,
                            mode       = 'mean',
                            upper_clip = 10,
                            maxiter    = 5,
                            maskmode   = (None, 'max')[n_flat>=3],
                            ncores     = ncores,
                            )
                #print('\033[{1}mCombined flat data for CCD {0}: \033[0m'.format(
                #    iccd+1, (34, 32, 31)[iccd]))
            flatdata_lst.append(flatdata)
            flatmask_lst.append(flatmask)

            # pack the combined flat data into flat_hdu_lst
            head = fits.Header()
            head['HIERARCH GAMSE FLAT CCD{} NFILE'.format(iccd+1)] = n_flat
            flat_hdu_lst.append(fits.ImageHDU(flatdata, head))
            flat_hdu_lst.append(fits.ImageHDU(flatmask))
        # CCD loop ends here

        # alias of flat data and mask
        flatdata1 = flatdata_lst[0].T
        flatmask1 = flatmask_lst[0].T
        flatdata2 = flatdata_lst[1].T
        flatmask2 = flatmask_lst[1].T
        flatdata3 = flatdata_lst[2].T
        flatmask3 = flatmask_lst[2].T

        # mosaic flat data
        flatdata, flatmask = mosaic_3_images(
                                data_lst = (flatdata1, flatdata2, flatdata3),
                                mask_lst = (flatmask1, flatmask2, flatmask3),
                                )

        flat_hdu_lst.append(fits.ImageHDU(flatdata.T))
        flat_hdu_lst.append(fits.ImageHDU(flatmask.T))
        # write flat data to file
        flat_hdu_lst = fits.HDUList(flat_hdu_lst)
        flat_hdu_lst.writeto(flat_file, overwrite=True)
        print('Flat data written to {}'.format(flat_file))

    ######################### find & trace orders ##########################

    # simple debackground for all 3 CCDs
    xnodes = np.arange(0, flatdata1.shape[1], 200)
    flatdbkg1 = simple_debackground(flatdata1, flatmask1, xnodes, smooth=20,
                    deg=3, maxiter=10)

    xnodes = np.arange(0, flatdata2.shape[1], 200)
    flatdbkg2 = simple_debackground(flatdata2, flatmask2, xnodes, smooth=20,
                    deg=3, maxiter=10)

    xnodes = np.arange(0, flatdata3.shape[1], 200)
    flatdbkg3 = simple_debackground(flatdata3, flatmask3, xnodes, smooth=20,
                    deg=3, maxiter=10)

    allimage, allmask = mosaic_3_images(
                        data_lst = (flatdbkg1, flatdbkg2, flatdbkg3),
                        mask_lst = (flatmask1, flatmask2, flatmask3),
                        )

    tracefig = TraceFigure()

    section = config['reduce.trace']
    aperset = find_apertures(allimage, allmask,
                scan_step  = section.getint('scan_step'),
                minimum    = section.getfloat('minimum'),
                separation = section.get('separation'),
                align_deg  = section.getint('align_deg'),
                filling    = section.getfloat('filling'),
                degree     = section.getint('degree'),
                display    = section.getboolean('display'),
                fig        = tracefig,
                )
    # decorate trace fig and save to file
    tracefig.adjust_positions()
    tracefig.suptitle('Trace for all 3 CCDs', fontsize=15)
    figfile = os.path.join(figpath, 'trace.png')
    tracefig.savefig(figfile)

    trcfile = os.path.join(midpath, 'trace.trc')
    aperset.save_txt(trcfile)

    regfile = os.path.join(midpath, 'trace.reg')
    aperset.save_reg(regfile, transpose=True)

    # save mosaiced flat image
    trace_hdu_lst = fits.HDUList(
                        [fits.PrimaryHDU(allimage.T),
                         fits.ImageHDU(allmask.T),
                        ])
    trace_hdu_lst.writeto(config['reduce.trace'].get('file'), overwrite=True)

    ######################### Extract flat spectrum ############################

    spectra1d = extract_aperset(flatdata, flatmask,
                    apertureset = aperset,
                    lower_limit = 6,
                    upper_limit = 6,
                    )

    flatmap = get_slit_flat(flatdata, flatmask,
                apertureset = aperset,
                spectra1d   = spectra1d,
                lower_limit = 6,
                upper_limit = 6,
                deg         = 7,
                q_threshold = 20**2,
                figfile     = 'spec_%02d.png',
                )
    fits.writeto('flat_resp.fits', flatmap, overwrite=True)

import numpy,import scipy,from scipy,import astropy


import os
import time
import random
import argparse
import numpy as np
import pandas as pd
import cv2
import PIL.Image
from tqdm import tqdm
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold
import torch
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data.sampler import RandomSampler, SequentialSampler
from torch.optim.lr_scheduler import CosineAnnealingLR
from util import GradualWarmupSchedulerV2
import apex
from apex import amp
from dataset import get_df, get_transforms, MelanomaDataset, SIIMISICDataset, enetv2
from models import Effnet_Melanoma, Resnest_Melanoma, Seresnext_Melanoma
from train import get_trans
import csv

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--kernel-type', type=str, required=True)
    parser.add_argument('--data-dir', type=str, default='/raid/')
    parser.add_argument('--data-folder', type=int, required=True)
    parser.add_argument('--image-size', type=int, required=True)
    parser.add_argument('--enet-type', type=str, required=True)
    parser.add_argument('--batch-size', type=int, default=64)
    parser.add_argument('--num-workers', type=int, default=32)
    parser.add_argument('--out-dim', type=int, default=9)
    parser.add_argument('--use-amp', action='store_true')
    parser.add_argument('--use-meta', action='store_true')
    parser.add_argument('--DEBUG', action='store_true')
    parser.add_argument('--model-dir', type=str, default='./weights')
    parser.add_argument('--log-dir', type=str, default='./logs')
    parser.add_argument('--sub-dir', type=str, default='./subs')
    parser.add_argument('--eval', type=str, choices=['best', 'best_20', 'final'], default="best")
    parser.add_argument('--n-test', type=int, default=8)
    parser.add_argument('--CUDA_VISIBLE_DEVICES', type=str, default='0')
    parser.add_argument('--n-meta-dim', type=str, default='512,128')
    parser.add_argument('--data-path', type=str, required=True)

    args, _ = parser.parse_known_args()
    return args

def main():

    df, df_test, meta_features, n_meta_features, mel_idx = get_df(
        args.kernel_type,
        args.out_dim,
        args.data_path,
        args.use_meta
    )

    transforms_train, transforms_val = get_transforms(args.image_size)

    if args.DEBUG:
        df_test = df_test.sample(args.batch_size * 3)
    dataset_test = MelanomaDataset(df_test, 'test', meta_features, transform=transforms_val)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=args.batch_size, num_workers=args.num_workers)

    # load model
    models = []
    for fold in range(5):

        if args.eval == 'best':
            model_file = os.path.join(args.model_dir, f'{args.kernel_type}_best_fold{fold}.pth')
        elif args.eval == 'best_20':
            model_file = os.path.join(args.model_dir, f'{args.kernel_type}_best_20_fold{fold}.pth')
        if args.eval == 'final':
            model_file = os.path.join(args.model_dir, f'{args.kernel_type}_final_fold{fold}.pth')

        model = ModelClass(
            args.enet_type,
            n_meta_features=n_meta_features,
            n_meta_dim=[int(nd) for nd in args.n_meta_dim.split(',')],
            out_dim=args.out_dim
        )
        model = model.to(device)

        try:  # single GPU model_file
            model.load_state_dict(torch.load(model_file), strict=True)
        except:  # multi GPU model_file
            state_dict = torch.load(model_file)
            state_dict = {k[7:] if k.startswith('module.') else k: state_dict[k] for k in state_dict.keys()}
            model.load_state_dict(state_dict, strict=True)

        if len(os.environ['CUDA_VISIBLE_DEVICES']) > 1:
            model = torch.nn.DataParallel(model)

        model.eval()
        models.append(model)

    # predict
    PROBS = []
    with torch.no_grad():
        for (data) in tqdm(test_loader):
            if args.use_meta:
                data, meta = data
                data, meta = data.to(device), meta.to(device)
                probs = torch.zeros((data.shape[0], args.out_dim)).to(device)
                for model in models:
                    for I in range(args.n_test):
                        l = model(get_trans(data, I), meta)
                        probs += l.softmax(1)
            else:
                data = data.to(device)
                probs = torch.zeros((data.shape[0], args.out_dim)).to(device)
                for model in models:
                    for I in range(args.n_test):
                        l = model(get_trans(data, I))
                        probs += l.softmax(1)

            probs /= args.n_test
            probs /= len(models)

            PROBS.append(probs.detach().cpu())

    PROBS = torch.cat(PROBS).numpy()

    # save cvs
    df_test['target'] = PROBS
    #df_test[['image_name', 'target']].to_csv(os.path.join(args.sub_dir, f'sub_{args.kernel_type}_{args.eval}.csv'), index=False)
    return df_test[['target']]

def infer_single_image(
    age,
    image_name,
    patient_id,
    sex,
    image_path,
    site,
    image_size,
    enet_type,
    kernel_type
    ):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    _, transforms_val = get_transforms(image_size)
    with open('datasettesting.csv', 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(["image_name", "patient_id", "sex", "age_approx", "anatom_site_general_challenge", "width", "height", "filepath"])
        writer.writerow([image_name, patient_id, sex, age, site, 6000, 4000, image_path])

    df_single_image = pd.read_csv('datasettesting.csv')

    dataset_test = SIIMISICDataset(df_single_image, 'test', 'test', transform=transforms_val)
    image = dataset_test[0]
    image = image.to(device).unsqueeze(0)

    models = []
    for i_fold in range(5):
        model = enetv2(enet_type, n_meta_features=0, out_dim=9)
        model = model.to(device)
        model_file = os.path.join(f'{kernel_type}_best_fold{i_fold}.pth')
        state_dict = torch.load(model_file)
        state_dict = {k.replace('module.', ''): state_dict[k] for k in state_dict.keys()}
        model.load_state_dict(state_dict, strict=True)
        model.eval()
        models.append(model)
    print(len(models))

    n_test = 8
    with torch.no_grad():
        probs = torch.zeros((image.shape[0], 9)).to(device)
        for model in models:
            for I in range(n_test):
                l = model(get_trans(image, I))
                probs += l.softmax(1)
        probs /= len(models) * n_test
    return probs

if __name__ == '__main__':
    """
    args = parse_args()
    os.makedirs(args.sub_dir, exist_ok=True)
    os.environ['CUDA_VISIBLE_DEVICES'] = args.CUDA_VISIBLE_DEVICES

    if args.enet_type == 'resnest101':
        ModelClass = Resnest_Melanoma
    elif args.enet_type == 'seresnext101':
        ModelClass = Seresnext_Melanoma
    elif 'efficientnet' in args.enet_type:
        ModelClass = Effnet_Melanoma
    else:
        raise NotImplementedError()

    DP = len(os.environ['CUDA_VISIBLE_DEVICES']) > 1

    device = torch.device('cuda')

    main()
    """
    age = 70
    image_name = "hello"
    patient_id = 'IP22423'
    sex = 'male'
    image_path = "testImage.jpg"
    site = "torso"
    image_size = 640
    enet_type = 'efficientnet-b7'
    kernel_type = "4c_b5ns_1.5e_640_ext_15ep"

    infer_single_image(
        age,
        image_name,
        patient_id,
        sex,
        image_path,
        site,
        image_size,
        enet_type,
        kernel_type
    )


#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import grpc
import json
import numpy as np
import os
import queue
import simulator_pb2
import simulator_pb2_grpc
import six
import time
import threading
from args import get_server_args
from collections import defaultdict
from concurrent import futures
from multi_head_ddpg import MultiHeadDDPG
from opensim_agent import OpenSimAgent
from opensim_model import OpenSimModel
from parl.utils import logger, ReplayMemory
from utils import calc_indicators, ScalarsManager, TransitionExperience

ACT_DIM = 19
VEL_DIM = 4
OBS_DIM = 185 + VEL_DIM
GAMMA = 0.96
TAU = 0.001
ACTOR_LR = 3e-5
CRITIC_LR = 3e-5
TRAIN_TIMES = 100
BATCH_SIZE = 128
NOISE_DECAY = 0.999998
_ONE_DAY_IN_SECONDS = 60 * 60 * 24

class SimulatorServer(simulator_pb2_grpc.SimulatorServicer):
    class ClientState(object):
        def __init__(self):
            self.memory = []  # list of Experience
            self.ident = None
            self.model_idx = np.random.randint(args.ensemble_num)
            self.last_target_changed = 0
            self.target_change_times = 0

        def reset(self):
            self.last_target_changed = 0
            self.memory = []
            self.model_idx = np.random.randint(args.ensemble_num)
            self.target_change_times = 0

        def update_last_target_changed(self):
            self.last_target_changed = len(self.memory)

    def __init__(self):
        self.rpm = ReplayMemory(int(2e6), OBS_DIM, ACT_DIM)

        # Need acquire lock when model learning or predicting
        self.locks = []
        for i in range(args.ensemble_num):
            self.locks.append(threading.Lock())

        models = []
        for i in range(args.ensemble_num):
            models.append(OpenSimModel(OBS_DIM, VEL_DIM, ACT_DIM, model_id=i))

        hyperparas = {
            'gamma': GAMMA,
            'tau': TAU,
            'ensemble_num': args.ensemble_num
        }
        alg = MultiHeadDDPG(models, hyperparas)

        self.agent = OpenSimAgent(alg, OBS_DIM, ACT_DIM, args.ensemble_num)

        self.scalars_manager = ScalarsManager(logger.get_dir())

        # add lock when appending data to rpm or writing scalars to tensorboard
        self.MEMORY_LOCK = threading.Lock()

        self.clients = defaultdict(self.ClientState)

        self.ready_client_queue = queue.Queue()

        self.noiselevel = 0.5
        self.global_step = 0

        # thread to keep training
        t = threading.Thread(target=self.keep_training)
        t.start()

    def _new_ready_client(self):
        """ The client is ready to start new episode,
        but blocking until training thread call client_ready_event.set()
    """
        client_ready_event = threading.Event()
        self.ready_client_queue.put(client_ready_event)
        logger.info(
            "[new_ready_client] approximate size of ready clients:{}".format(
                self.ready_client_queue.qsize()))
        client_ready_event.wait()

    def Send(self, request, context):
        """ Implement Send function in SimulatorServicer
        Everytime a request comming, will create a new thread to handle
    """
        ident, obs, reward, done, info = request.id, request.observation, request.reward, request.done, request.info
        client = self.clients[ident]
        info = json.loads(info)

        if 'first' in info:
            # Waiting training thread to allow start new episode
            self._new_ready_client()

        obs = np.array(obs, dtype=np.float32)
        self._process_msg(ident, obs, reward, done, info)

        if done:
            # Waiting training thread to allow start new episode
            self._new_ready_client()

        action = self.pred_batch(obs, client.model_idx)
        step = len(client.memory) - client.last_target_changed

        # whether to add noise depends on the ensemble_num
        if args.ensemble_num == 1:
            current_noise = self.noiselevel * (0.98**(step - 1))
            noise = np.zeros((ACT_DIM, ), dtype=np.float32)
            if ident % 3 == 0:
                if step % 5 == 0:
                    noise = np.random.randn(ACT_DIM) * current_noise
            elif ident % 3 == 1:
                if step % 5 == 0:
                    noise = np.random.randn(ACT_DIM) * current_noise * 2
            action += noise
        action = np.clip(action, -1, 1)
        client.memory[-1].action = action
        extra_info = {}
        return simulator_pb2.Reply(action=action, extra=json.dumps(extra_info))

    def _process_msg(self, ident, obs, reward, done, info):
        client = self.clients[ident]
        reward_scale = (1 - GAMMA)
        info['shaping_reward'] *= reward_scale
        if len(client.memory) > 0:
            client.memory[-1].reward = reward
            info['target_change_times'] = client.target_change_times
            client.memory[-1].info = info
            if info['target_changed']:
                client.target_change_times = min(
                    client.target_change_times + 1, 3)
                # re-sample model_idx after target was changed
                client.model_idx = np.random.randint(args.ensemble_num)
            if done:
                assert 'last_obs' in info
                self._parse_memory(client, ident, info['last_obs'])
        client.memory.append(
            TransitionExperience(obs=obs, action=None, reward=None, info=None))
        if 'target_changed' in info and info['target_changed']:
            client.update_last_target_changed()
        return False

    def _parse_memory(self, client, ident, last_obs):
        mem = client.memory
        n = len(mem)

        # debug info
        if ident == 1:
            for i, exp in enumerate(mem):
                logger.info(
                    "[step:{}] obs:{} action:{} reward:{} shaping_reward:{}".
                    format(i, np.sum(mem[i].obs), np.sum(mem[i].action),
                           mem[i].reward, mem[i].info['shaping_reward']))

        episode_rpm = []
        for i in range(n - 1):
            if not mem[i].info['target_changed']:
                episode_rpm.append([
                    mem[i].obs, mem[i].action, mem[i].info['shaping_reward'],
                    mem[i + 1].obs, False, mem[i].info['target_change_times']
                ])
        if not mem[-1].info['target_changed']:
            episode_rpm.append([
                mem[-1].obs, mem[-1].action, mem[-1].info['shaping_reward'],
                last_obs, not mem[-1].info['timeout'],
                mem[i].info['target_change_times']
            ])

        indicators_dict = calc_indicators(mem)
        indicators_dict['free_client_num'] = self.ready_client_queue.qsize()
        indicators_dict['noiselevel'] = self.noiselevel

        with self.MEMORY_LOCK:
            self.add_episode_rpm(episode_rpm)
            self.scalars_manager.record(indicators_dict, self.global_step)
            self.global_step += 1
            if self.global_step >= 50:
                self.noiselevel = self.noiselevel * NOISE_DECAY

        client.reset()

    def learn(self):
        result_q = queue.Queue()
        th_list = []
        for j in range(args.ensemble_num):
            t = threading.Thread(
                target=self.train_single_model, args=(j, result_q))
            th_list.append(t)
        start_time = time.time()
        for t in th_list:
            t.start()
        for t in th_list:
            t.join()

        logger.info("[learn] {} heads, time consuming:{}".format(
            args.ensemble_num,
            time.time() - start_time))
        for t in th_list:
            result = result_q.get()
            for critic_loss in result:
                self.scalars_manager.feed_critic_loss(critic_loss)

    def train_single_model(self, model_idx, result_q):
        logger.info("[train_single_model] model_idx:{}".format(model_idx))
        critic_loss_list = []
        lock = self.locks[model_idx]
        memory = self.rpm

        actor_lr = ACTOR_LR * (1.0 - 0.05 * model_idx)
        critic_lr = CRITIC_LR * (1.0 + 0.1 * model_idx)

        for T in range(TRAIN_TIMES):
            [states, actions, rewards, new_states,
             dones] = memory.sample_batch(BATCH_SIZE)
            lock.acquire()
            critic_loss = self.agent.learn(states, actions, rewards,
                                           new_states, dones, actor_lr,
                                           critic_lr, model_idx)
            lock.release()
            critic_loss_list.append(critic_loss)
        result_q.put(critic_loss_list)

    def keep_training(self):
        episode_count = 1000000
        for T in range(episode_count):
            if self.rpm.size() > BATCH_SIZE * args.warm_start_batchs:
                self.learn()
                logger.info(
                    "[keep_training/{}] trying to acq a new env".format(T))

            # Keep training and predicting balance
            # After training, waiting for a ready client, and set the client start new episode
            ready_client_event = self.ready_client_queue.get()
            ready_client_event.set()

            if np.mod(T, 100) == 0:
                logger.info("saving models")
                self.save(T)
            if np.mod(T, 10000) == 0:
                logger.info("saving rpm")
                self.save_rpm()

    def save_rpm(self):
        save_path = os.path.join(logger.get_dir(), "rpm.npz")
        self.rpm.save(save_path)

    def restore_rpm(self, rpm_dir):
        self.rpm.load(rpm_dir)

    def save(self, T):
        save_path = os.path.join(logger.get_dir(),
                                 'model_every_100_episodes/step-{}'.format(T))
        self.agent.save_params(save_path)

    def restore(self, model_path, restore_from_one_head):
        logger.info('restore model from {}'.format(model_path))
        self.agent.load_params(model_path, restore_from_one_head)

    def add_episode_rpm(self, episode_rpm):
        for x in episode_rpm:
            self.rpm.append(
                obs=x[0], act=x[1], reward=x[2], next_obs=x[3], terminal=x[4])

    def pred_batch(self, obs, model_idx=None):
        assert model_idx is not None
        batch_obs = np.expand_dims(obs, axis=0)
        self.locks[model_idx].acquire()
        action = self.agent.predict(batch_obs, model_idx)
        self.locks[model_idx].release()
        action = np.squeeze(action, axis=0)
        return action

class SimulatorHandler(threading.Thread):
    def __init__(self, simulator_server):
        threading.Thread.__init__(self)
        self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=400))
        simulator_pb2_grpc.add_SimulatorServicer_to_server(
            simulator_server, self.server)
        self.server.add_insecure_port('[::]:{}'.format(args.port))

    def run(self):
        self.server.start()
        try:
            while True:
                time.sleep(_ONE_DAY_IN_SECONDS)
        except KeyboardInterrupt:
            self.server.stop(0)

if __name__ == '__main__':
    args = get_server_args()

    if args.logdir is not None:
        logger.set_dir(args.logdir)

    simulator_server = SimulatorServer()

    if args.restore_rpm_path is not None:
        simulator_server.restore_rpm(args.restore_rpm_path)
    if args.restore_model_path is not None:
        simulator_server.restore(args.restore_model_path,
                                 args.restore_from_one_head)

    simulator_hanlder = SimulatorHandler(simulator_server=simulator_server)
    simulator_hanlder.run()


import numpy as np
import torch

from multiworld.envs.mujoco.cameras import sawyer_init_camera_multiobject_push
from rlkit.launchers.launcher_util import run_experiment
from rlkit.launchers.skewfit_experiments import skewfit_full_experiment
from rlkit.torch import pytorch_util as ptu
from rlkit.torch.vae.conv_vae import imsize48_architecture_2
import rlkit.torch.vae.vae_schedules as vae_schedules

from path_length_settings import get_path_settings

experiment = skewfit_full_experiment

path_settings = get_path_settings('Skewfit', 'Rearrange', n_objects=1)
PATH_LENGTH = path_settings.train_path_length
PATH_LENGTH_EVAL = path_settings.eval_path_length

vae_architecture = imsize48_architecture_2.copy()

variant = dict(
    algorithm='Skew-Fit',
    double_algo=False,
    online_vae_exploration=False,
    imsize=48,
    init_camera=sawyer_init_camera_multiobject_push,
    env_id='SawyerMultiobjectRearrangeEnv-OneObj-v0',
    skewfit_variant=dict(
        save_video=True,
        custom_goal_sampler='replay_buffer',
        online_vae_trainer_kwargs=dict(
            beta=0.2,
            lr=1e-3,
            batch_size=128,
            img_dump_interval=50
        ),
        save_video_period=100,
        qf_kwargs=dict(
            hidden_sizes=[400, 300]
        ),
        policy_kwargs=dict(
            hidden_sizes=[400, 300]
        ),
        max_path_length=PATH_LENGTH,
        max_path_length_eval=PATH_LENGTH_EVAL,
        algo_kwargs=dict(
            batch_size=2048,
            num_epochs=1201,
            num_eval_steps_per_epoch=1 * PATH_LENGTH_EVAL,
            num_train_loops_per_epoch=1,
            num_expl_steps_per_train_loop=10 * PATH_LENGTH,
            num_trains_per_train_loop=20 * PATH_LENGTH,
            min_num_steps_before_training=10000,
            vae_training_schedule=vae_schedules.custom_schedule_3,
            oracle_data=False,
            vae_save_period=25,
            parallel_vae_train=False,
        ),
        twin_sac_trainer_kwargs=dict(
            policy_lr=1e-3,
            qf_lr=1e-3,
            discount=0.95,
            reward_scale=1,
            soft_target_tau=1e-2,
            target_update_period=1,
            use_automatic_entropy_tuning=True,
        ),
        replay_buffer_kwargs=dict(
            start_skew_epoch=10,
            max_size=100000,
            fraction_goals_rollout_goals=0.2,
            fraction_goals_env_goals=0.5,
            exploration_rewards_type='None',
            vae_priority_type='vae_prob',
            priority_function_kwargs=dict(
                sampling_method='importance_sampling',
                decoder_distribution='gaussian_identity_variance',
                num_latents_to_sample=10,
            ),
            power=-0.5,
            relabeling_goal_sampling_mode='custom_goal_sampler',
            refresh_latents_batch_size=512
        ),
        exploration_goal_sampling_mode='custom_goal_sampler',
        evaluation_goal_sampling_mode='reset_of_env',
        normalize=False,
        render=False,
        exploration_noise=0.0,
        exploration_type='ou',
        training_mode='train',
        testing_mode='test',
        reward_params=dict(
            type='latent_distance',
        ),
        observation_key='latent_observation',
        desired_goal_key='latent_desired_goal',
        vae_wrapped_env_kwargs=dict(
            sample_from_true_prior=True,
        ),
    ),
    train_vae_variant=dict(
        representation_size=4,
        beta=0.2,
        num_epochs=0,
        dump_skew_debug_plots=False,
        decoder_activation='gaussian',
        generate_vae_dataset_kwargs=dict(
            N=50,
            test_p=.1,
            use_cached=False,
            show=False,
            oracle_dataset=True,
            oracle_dataset_using_set_to_goal=True,
            n_random_steps=100,
            non_presampled_goal_img_is_garbage=True,
        ),
        vae_kwargs=dict(
            input_channels=3,
            architecture=vae_architecture,
            decoder_distribution='gaussian_identity_variance',
            hidden_init=torch.nn.init.orthogonal_,
            mse_reduction='sum',
            normalize_logprob=True
        ),
        algo_kwargs=dict(
            start_skew_epoch=5000,
            is_auto_encoder=False,
            batch_size=64,
            lr=1e-3,
            skew_config=dict(
                method='vae_prob',
                power=-1,
            ),
            skew_dataset=True,
            priority_function_kwargs=dict(
                decoder_distribution='gaussian_identity_variance',
                sampling_method='importance_sampling',
                num_latents_to_sample=10,
            ),
            use_parallel_dataloading=False,
            img_dump_interval=25
        ),

        save_period=25,
    ),
)

if __name__ == "__main__":
    exp_prefix = 'skewfit-rearrange-1-obj'

    run_experiment(
        skewfit_full_experiment,
        exp_prefix=exp_prefix,
        mode='local',
        variant=variant,
        use_gpu=True)


import numpy as np
from tqdm import tqdm
from typing import List

from intuit_topic.backend import BaseEmbedder

class USEBackend(BaseEmbedder):
    """ Universal Sentence Encoder

    USE encodes text into high-dimensional vectors that
    are used for semantic similarity in BERTopic.

    Arguments:
        embedding_model: An USE embedding model

    Usage:

    ```python
    import tensorflow_hub
    from intuit_topic.backend import USEBackend

    embedding_model = tensorflow_hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
    use_embedder = USEBackend(embedding_model)
    ```
    """
    def __init__(self, embedding_model):
        super().__init__()

        try:
            embedding_model(["test sentence"])
            self.embedding_model = embedding_model
        except TypeError:
            raise ValueError("Please select a correct USE model: \n"
                             "`import tensorflow_hub` \n"
                             "`embedding_model = tensorflow_hub.load(path_to_model)`")

    def embed(self,
              documents: List[str],
              verbose: bool = False) -> np.ndarray:
        """ Embed a list of n documents/words into an n-dimensional
        matrix of embeddings

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        embeddings = np.array([self.embedding_model([doc]).cpu().numpy()[0]
                               for doc in tqdm(documents, disable=not verbose)])
        return embeddings


"""NIRC2 specific methods and variables for an AOInstrument.
"""
from __future__ import division, print_function
import astropy.io.fits as pyfits
import numpy as np
import scipy.ndimage as nd
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import time
import glob
import pdb
import time

from .aoinstrument import AOInstrument

class NIRC2(AOInstrument):
 """The NIRC2 Class, that enables processing of NIRC2 images.
 """
 #A global definition, for error-checking downstream
 instrument = 'NIRC2'
 def is_bad_surrounded(self,bad):
    #Returns matrix of booleans showing which pixels are surrounded by bad pixels
    #"Surrounded" means there is at least one bad pixel in at least two directions
    numPixels = 3
    sz = len(bad)
    is_bad_to_left = np.zeros((sz,sz-numPixels))
    is_bad_to_right = np.zeros((sz,sz-numPixels))
    is_bad_above = np.zeros((sz-numPixels,sz))
    is_bad_below = np.zeros((sz-numPixels,sz))
    for ii in range(0,numPixels):
        is_bad_to_left+=bad[0:sz,numPixels-ii-1:sz-ii-1]
        is_bad_to_right+=bad[0:sz,ii+1:sz-numPixels+ii+1]
        is_bad_above+=bad[numPixels-ii-1:sz-ii-1,0:sz]
        is_bad_below+=bad[ii+1:sz-numPixels+ii+1,0:sz]
    is_bad_to_left = is_bad_to_left>0
    is_bad_to_right = is_bad_to_right>0
    is_bad_above = is_bad_above>0
    is_bad_below = is_bad_below>0
    is_surrounded = np.zeros((sz,sz))
    is_surrounded[0:sz,numPixels:sz]+=is_bad_to_left
    is_surrounded[0:sz,0:sz-numPixels]+=is_bad_to_right
    is_surrounded[numPixels:sz,0:sz]+=is_bad_above
    is_surrounded[0:sz-numPixels,0:sz]+=is_bad_below
    is_surrounded = is_surrounded>2
    return is_surrounded

 def saturated_pixels(self, image, header, threshold=17500):
    """Returns coordinates of all saturated pixels
    Uses image (already corrected for nonlinearity)
    and header from file

    Parameters
    ----------
    image: numpy array
        The input image

    header: pyfits header
        The header from this image.
    """
    if "COADDS" in header.keys():
        pixels = np.where(image/header["COADDS"]>threshold)
    else:
        pixels = np.where(image>threshold)
    return pixels

 def make_all_darks(self, ddir='', rdir=''):
    """Make all darks in a current directory. This skeleton routine assumes that
    keywords "SHRNAME", "NAXIS1" and "NAXIS2" exist.
    """
    #Allow over-riding default reduction and data directories.
    if (rdir == ''):
        rdir = self.rdir
    if (ddir == ''):
        ddir = self.ddir
    if len(self.csv_dict) == 0:
        print("Error: Run read_summary_csv first. No darks made.")
        return
    darks = np.where(np.array(self.csv_dict['SHRNAME']) == 'closed')[0]
    #Now we need to find unique values of the following:
    #NAXIS1, NAXIS2 (plus for nirc2... ITIME, COADDS, MULTISAM)
    codes = []
    for d in darks:
        codes.append(self.csv_dict['NAXIS1'][d] + self.csv_dict['NAXIS2'][d] +
            self.csv_dict['ITIME'][d] + self.csv_dict['COADDS'][d] + self.csv_dict['MULTISAM'][d])
    codes = np.array(codes)
     #For each unique code, find all dark files and call make_dark.
    for c in np.unique(codes):
        w = np.where(codes == c)[0]
        if (len(w) >= 3):
            files = [ddir + self.csv_dict['FILENAME'][darks[ww]] for ww in w]
            self.make_dark(files, rdir=rdir)

 def make_all_flats(self, ddir='', rdir=''):
    """Search for sets of files that look like they are a series of flats. If "Lamp Off"
    files exist within 100 files or so of the flats, call them the darks to go with the
    flats. """
    #Allow over-riding default reduction and data directories.
    if (rdir == ''):
        rdir = self.rdir
    if (ddir == ''):
        ddir = self.ddir
    if len(self.csv_dict) == 0:
        print("Error: Run read_summary_csv first. No darks made.")
        return
    #Fill in elevation with a default value (45, for dome flat position) if there are fits header errors.
    els = self.csv_dict['EL']
    for i in range(len(els)):
        try:
            this_el = float(els[i])
        except:
            els[i] = '45.0'
    els = els.astype(float)
    #If we're in the dome flat position with more than 1000 counts, this looks
    #like it could be a dome flat!
    flats_maybe = np.where( (self.csv_dict['MEDIAN_VALUE'].astype('float') > 1000) &
                            (np.abs(els - 45) < 0.01) )[0]
    codes = []
    fluxes = self.csv_dict['MEDIAN_VALUE'][flats_maybe].astype(float)
    for ix in range(len(els)):
        codes.append(self.csv_dict['FILTER'][ix] + self.csv_dict['NAXIS1'][ix] + self.csv_dict['NAXIS2'][ix] +
            self.csv_dict['ITIME'][ix] + self.csv_dict['COADDS'][ix] + self.csv_dict['MULTISAM'][ix] +
            self.csv_dict['SLITNAME'][ix])
    codes = np.array(codes)
    flat_codes = codes[flats_maybe]
    #For each unique code, find the files with consistent flux
    for c in np.unique(flat_codes):
        #w indexes flats_maybe
        w = np.where(flat_codes == c)[0]
        #Flux has to be within 10% of the median to count.
        this_flat_flux = np.median(fluxes[w])
        good_flats = flats_maybe[w[np.where(np.abs( (fluxes[w] - this_flat_flux)/this_flat_flux < 0.1))[0]]]
        #Less than 5 flats... don't bother.
        if (len(good_flats) >= 5):
            ffiles = [ddir + self.csv_dict['FILENAME'][ww] for ww in good_flats]
            lamp_off = np.where( (codes == c) & (np.array(self.csv_dict['MEDIAN_VALUE'].astype(float) < 600) & \
                    (np.abs(els - 45) < 0.01) ) )[0]
            if (len(lamp_off) >= 3):
                #Use these lamp_off indexes to create a "special" dark.
                dfiles = [ddir + self.csv_dict['FILENAME'][ww] for ww in lamp_off]
                try:
                    hh = pyfits.open(dfiles[0], ignore_missing_end=True)[0].header
                except:
                    hh = pyfits.open(dfiles[0]+'.gz', ignore_missing_end=True)[0].header
                dfilename = str(lamp_off[0]) + '_' + self.get_dark_filename(hh)
                self.make_dark(dfiles, out_file=dfilename)
                self.make_flat(ffiles, dark_file=dfilename)
            #Otherwise, just use default darks. This *will* give an error if they don't exist.
            else:
                self.make_flat(ffiles)

 def csv_block_string(self, ix):
    """Find a string from the summary csv file that identifies a unique configuration
    for a set of files to be processed as a block. It isn't *quite* correct
    because the target name sometimes stays the same with a target change.

    Parameters
    ----------
    ix: int
        The index of the file (in the csv dictionary) that we want to get a block string
        for"""
    if len(self.csv_dict) == 0:
        print("Error: Run read_summary_csv first. No string returned.")
        return
    block_string = self.csv_dict['NAXIS1'][ix] + self.csv_dict['NAXIS2'][ix] + \
        self.csv_dict['TARGNAME'][ix] + self.csv_dict['FILTER'][ix] + \
        self.csv_dict['ITIME'][ix] + self.csv_dict['COADDS'][ix]
    return block_string

 def info_from_header(self, h, subarr=None):
    """Find important information from the fits header and store in a common format

    Parameters
    ----------
    h:    The fits header

    Returns
    -------
    (dark_file, flat_file, filter, wave, rad_pixel)
    """

    #First, sanity check the header
    try: inst=h['CURRINST']
    except: inst=''
    if (len(inst)==0):
        print("Error: could not find instrument in header...")
        raise UserWarning
    if ((self.instrument != inst) & (inst[0:3] != '###')):
        print("Error: software expecting: ", self.instrument, " but instrument is: ", inst)
        raise UserWarning

    try: fwo = h['FWONAME']
    except:
        print("No FWONAME in NIRC2 header")
        raise UserWarning
    try: fwi = h['FWINAME']
    except:
        print("No FWINAME in NIRC2 header")
        raise UserWarning
    try: slit = h['SLITNAME']
    except:
        print("No SLITNAME in NIRC2 header")
        raise UserWarning

    if (fwi=='Kp'):
        wave = 2.12e-6
        filter='Kp'
    elif (fwi=='CH4_short'):
        wave = 1.60e-6
        filter='CH4_short'
    elif (fwo=='Kcont'):
        wave = 2.27e-6
        filter='Kcont'
    elif (fwi=='Lp'):
        wave = 3.776e-6
        filter='Lp'
    elif (fwi=='H2O_ice'):
        wave = 3.1e-6
        filter='H2O_ice'
    elif (fwo=='PAH'):
        wave = 3.3e-6
        filter='PAH'
    elif (fwi=='Ms'):
        #The central wavelength for this filter in the manual is 4.67 microns, with
        #a width of 0.24 microns. We will set the wavelength to the long edge of this
        #wavelength range in order to ensure we have enough signal - data tend to be
        #noisy in this filter!
        wave = 4.79e-6
        filter = 'Ms'
    elif (fwi=='CH4_long'):
        wave   = 1.68e-6
        filter = 'CH4_long'
    elif (fwi=='J'):
        wave = 1.248e-6
        filter = 'J'
    elif (fwi=='z'):
        wave   = 1.03e-6
        filter = 'z'
    elif (fwi=='Y'):
        wave = 1.018e-6
        filter = 'Y'
    elif (fwi=='H'):
        wave=1.633e-6
        filter='H'
    elif (fwo=='Hcont'):
        wave=1.5804e-6
        filter='Hcont'
    elif (fwo=='Br_alpha'):
        wave=4.052e-6
        filter='Br_alpha'
    elif (fwo=='Bra_cont'):
        wave=3.987e-6
        filter='Bra_cont'
    else:
        print("Unknown Filter!")
        pdb.set_trace()
    if (slit == 'none'):
        flat_file = 'flat_' + filter + '.fits'
    else:
        flat_file = 'flat_' + filter + '_' + slit + '.fits'

    try: camname = h['CAMNAME']
    except:
           print("No CAMNAME in header")
    if (camname == 'narrow'):
        #This comes from the Yelda (2010) paper.
         rad_pixel = 0.009952*(np.pi/180.0/3600.0)
    else:
        print("Unknown Camera!")
        raise UserWarning
    #Estimate the expected readout noise directly from the header.
    if h['SAMPMODE'] == 2:
        multisam = 1
    else:
        multisam = h['MULTISAM']
    #The next line comes from the NIRC2 manual home page.
    gain = 4.0
    rnoise = 50.0/gain/np.sqrt(multisam)*np.sqrt(h['COADDS'])
    #Find the appropriate dark file if needed.
    dark_file = self.get_dark_filename(h)
    if ( (h['OBSDNAME'] == 'sodiumDichroic') & (h['TARGNAME'][0:2] == 'tt')):
        targname = h['OBJECT']
    else:
        targname = h['TARGNAME']
    #The pupil orientation...
    try:
        el = float(h['EL'])
    except:
        el = -1
    if (el > 0):
        vertang_pa = h['ROTPPOSN']-h['EL']-h['INSTANGL']
        pa = vertang_pa + h['PARANG']
    else:
        vertang_pa=np.NaN
        pa = np.NaN
    #Find the pupil type and parameters for the pupil...
    pupil_params=dict()
    if (fwo == '9holeMsk'):
        pupil_type = 'circ_nrm'
        hole_xy = [[3.44,  4.57,  2.01,  0.20, -1.42, -3.19, -3.65, -3.15,  1.18],
                   [-2.22, -1.00,  2.52,  4.09,  4.46,  0.48, -1.87, -3.46, -3.01]]
        #For some reason... the y-axis is reversed, and we will transpose here
        #to save any transposes later.
        #NB as we aren't doing (u,v) coordinates but only chip coordinates here,
        #there is no difference between reversing the x- and y- axes.
        hole_xy = np.array(hole_xy)
        hole_xy[1,:] = -hole_xy[1,:]
        hole_xy = hole_xy[::-1,:]
        pupil_params['hole_xy'] = hole_xy
        pupil_params['hole_diam'] = 1.1
        pupil_params['mask_rotation'] = -0.01
        pupil_params['mask'] = 'g9'
        if subarr:
            subarr_string = '_' + str(subarr)
        else:
            subarr_string = ''
        fourier_mask_file = 'fmask_' + filter + '_'+ pupil_params['mask'] + subarr_string + '.fits'
        kernel_defn_file = 'kernel_' + filter + '_'+ pupil_params['mask'] + subarr_string + '.fits'
    elif (fwi == '18holeMsk'):
        pupil_type = 'circ_nrm'
        pupil_params['mask'] = 'g18'
        print("Still to figure out 18 hole mask...")
        raise UserWarning
    else:
        pupil_type = 'keck'
        pupil_params['segment_size'] = 1.558
        pupil_params['obstruction_size'] = 2.0 #Guessed
        fourier_mask_file = 'fmask_' + filter + '_fullpupil.fits'
        kernel_defn_file = 'kernel_' + filter + '_fullpupil.fits'
#    else:
#        print "Assuming full pupil..."
#        pupil_type = 'annulus'
#        pupil_params['inner_diam'] = 1.8
#        pupil_params['outer_diam'] = 10.2 #Maximum diameter is really 10.5
#        ftpix_file = 'ftpix_' + filter + '_fullpupil.fits'
    return {'dark_file':dark_file, 'flat_file':flat_file, 'filter':filter,
        'wave':wave, 'rad_pixel':rad_pixel,'targname':targname,
        'pupil_type':pupil_type,'pupil_params':pupil_params,
        'fourier_mask_file':fourier_mask_file, 'kernel_defn_file':kernel_defn_file,
        'gain':gain, 'rnoise':rnoise, 'vertang_pa':vertang_pa, 'pa':pa}

 def get_dark_filename(self,h):
    """Create a dark fits filename based on a header

    Parameters
    ----------
    h: header from astropy.io.fits

    Returns
    -------
    dark_file: string
    """
    if h['SAMPMODE'] == 2:
        multisam = 1
    else:
        multisam = h['MULTISAM']
    dark_file = 'dark_' + str(h['NAXIS1']) +'_'+str(h['COADDS']) +'_' +str(multisam)+'_'+ str(int(h['ITIME']*100)) + '.fits'
    return dark_file

 def destripe_nirc2(self,im, subtract_edge=True, subtract_median=False, do_destripe=True):
    """Destripe an image from the NIRC2 camera.

    The algorithm is:
    1) Subtract the mode from each quadrant.
    2) For each pixel, find the 24 pixels in other quadrants corresponding to
    reflections about the chip centre.
    3) Subtract the median of these pixels.

    Parameters
    ----------
    im: array_like
        The input image.
    subtract_median: bool, optional
        Whether or not to subtract the median from each quadrant.
    subtract_edge: bool, optional
        Whether or not to adjust the means of each quadrant by the edge pixels.

    Returns
    -------
    im: array_like
        The corrected image.
    """
    s = im.shape
    quads = [im[0:s[0]//2,0:s[1]//2],im[s[0]:s[0]//2-1:-1,0:s[1]//2],
             im[0:s[0]//2,s[1]:s[1]//2-1:-1],im[s[0]:s[0]//2-1:-1,s[1]:s[1]//2-1:-1]]
    quads = np.array(quads, dtype='float')
    #Work through the quadrants, modifying based on the edges.
    if subtract_edge:
        quads[1] += np.median(quads[3][:,s[1]//2-8:s[1]//2])- np.median(quads[1][:,s[1]//2-8:s[1]//2])
        quads[2] += np.median(quads[3][s[0]//2-8:s[0]//2,:])- np.median(quads[2][s[0]//2-8:s[0]//2,:])
        delta = 0.5*(np.median(quads[3][s[0]//2-8:s[0]//2,:]) + np.median(quads[3][:,s[1]//2-8:s[1]//2])
               - np.median(quads[0][s[0]//2-8:s[0]//2,:]) - np.median(quads[0][:,s[1]//2-8:s[1]//2]))
        quads[0] += delta
    #Subtract the background
    if subtract_median:
        print("Subtracting Medians...")
        MED_DIFF_MULTIPLIER = 4.0
        for i in range(4):
            quad = quads[i,:,:]
            med = np.median(quad)
            dispersion = np.median(np.abs(quad - med))
            goodpix = np.where(np.abs(quad - med) < MED_DIFF_MULTIPLIER*dispersion)
            med = np.median(quad[goodpix])
            quads[i,:,:] -= med
    if do_destripe:
        quads = quads.reshape((4,s[0]//2,s[1]//16,8))
        stripes = quads.copy()
        for i in range(4):
            for j in range(s[0]//2): #The -1 on  line is because of artifacts
                for k in range(s[0]//16):
                    pix = np.array([stripes[(i+1)%4,j,k,:],stripes[(i+2)%4,j,k,:],stripes[(i+3)%4,j,k,:]])
                    quads[i,j,k,:] -= np.median(pix)
        quads = quads.reshape((4,s[0]//2,s[1]//2))
    im[0:s[0]//2,0:s[1]//2] = quads[0]
    im[s[0]:s[0]//2-1:-1,0:s[1]//2] = quads[1]
    im[0:s[0]//2,s[1]:s[1]//2-1:-1] = quads[2]
    im[s[0]:s[0]//2-1:-1,s[1]:s[1]//2-1:-1] = quads[3]
    return im

 def make_dark(self,in_files, out_file='', subtract_median=True, destripe=True, med_threshold=15.0, rdir=''):
    """Create a dark frame and save to a fits file,
    with an attached bad pixel map as the first fits extension.

    Parameters
    ----------
    in_files : array_like (dtype=string). A list if input filenames.
    out_file: string
        The file to write to.
    subtract_median: bool, optional
        Whether or not to subtract the median from each frame (or quadrants)
    destripe: bool, optional
        Whether or not to destripe the images.
    med_threshold: float, optional
        The threshold for pixels to be considered bad if their absolute
        value differs by more than this multiple of the median difference
        of pixel values from the median.

    Returns
    -------
    (optional) out_file: If an empty string is given, it is filled with the default out
    filename
    """
    #Allow over-riding default reduction directory.
    if (rdir == ''):
        rdir = self.rdir
    VAR_THRESHOLD = 10.0
    nf = len(in_files)
    if (nf < 3):
        print("At least 3 dark files sre needed for reliable statistics")
        raise UserWarning
    # Read in the first dark to check the dimensions.
    try:
        in_fits = pyfits.open(in_files[0], ignore_missing_end=True)
    except:
        in_fits = pyfits.open(in_files[0]+'.gz', ignore_missing_end=True)
    h = in_fits[0].header
    instname = ''
    try: instname=h['CURRINST']
    except:
        print("Unknown Header Type")
    #Create the output filename if needed
    if (out_file == ''):
        out_file = self.get_dark_filename(h)
    s = in_fits[0].data.shape
    in_fits.close()
    darks = np.zeros((nf,s[0],s[1]))
    plt.clf()
    for i in range(nf):
        #Read in the data, linearizing as a matter of principle, and also because
        #this routine is used for
        adark = self.linearize_nirc2(in_files[i])
        if (instname == 'NIRC2'):
            adark = self.destripe_nirc2(adark, subtract_median=subtract_median, do_destripe=destripe)
        if (subtract_median):
            plt.imshow(np.minimum(adark,1e2))
        else:
            plt.imshow(adark)
            print("Median: " + str(np.median(adark)))
        plt.pause(0.001)
        #plt.draw()
        darks[i,:,:] = adark
    #Now look for weird pixels.
    med_dark = np.median(darks, axis=0)
    max_dark = np.max(darks, axis=0)
    var_dark = np.zeros((s[0],s[1]))
    for i in range(nf):
        var_dark += (darks[i,:,:] - med_dark)**2
    var_dark -= (max_dark - med_dark)**2
    var_dark /= nf-2
    #We need to threshold the med_diff quantity in case of low-noise, many subread images
    med_diff = np.maximum(np.median(np.abs(med_dark - np.median(med_dark))),1.0)
    print("Median difference: " + str(med_diff))
    med_var_diff = np.median(np.abs(var_dark - np.median(var_dark)))
    bad_med = np.abs(med_dark - np.median(med_dark)) > med_threshold*med_diff
    bad_var = np.abs(var_dark) > np.median(var_dark) + VAR_THRESHOLD*med_var_diff
    print("Pixels with bad mean: " + str(np.sum(bad_med)))
    print("Pixels with bad var: " + str(np.sum(bad_var)))
    bad = np.logical_or(bad_med, bad_var)
    med_dark[bad] = 0.0
    #Copy the original header to the dark file.
    hl = pyfits.HDUList()
    hl.append(pyfits.ImageHDU(med_dark,h))
    hl.append(pyfits.ImageHDU(np.uint8(bad)))
    hl.writeto(rdir+out_file,clobber=True)
    plt.figure(1)
    plt.imshow(med_dark,cmap=cm.gray, interpolation='nearest')
    plt.title('Median Frame')
    plt.figure(2)
    plt.imshow(bad,cmap=cm.gray, interpolation='nearest')
    plt.title('Bad Pixels')
    plt.pause(0.001)
    #plt.draw()

 def linearize_nirc2(self,in_file, out_file=''):
    """Procedure to linearize NIRC2 (treated as a single detector).
    Run on all images nominally before running anything else.
    Originally from IDL code by Stan Metchev with Adam Kraus modifications.

    Parameters
    ----------
    in_file: string
        The input fits file
    out_file: string, optional
        An output fits file

    Returns
    -------
    im: (N,N) array
        The linearized image.

    Notes
    -----
    This procedure takes COADDS into account, but does not take subreads into account.
    i.e. in MCDS sampling, there is slightly more nonlinearity than accounted for,
    because the detector is more saturated by the time the last read is made.
    """

    coeff = np.array([1.001,-6.9e-6,-0.70e-10])
    #Get key parameters from the header and get the data
    try:
        in_fits = pyfits.open(in_file, ignore_missing_end=True)
    except:
        in_fits = pyfits.open(in_file + '.gz', ignore_missing_end=True)
    z = in_fits[0].header
    fitsarr = in_fits[0].data
    in_fits.close()
    #See if we've already updated this header
    try:
        lindate = z['LINHIST']
    except:
        print('Linearizing: ' + in_file)
        xsub=z['NAXIS1']
        ysub=z['NAXIS2']
        norm=np.array((xsub,ysub))
        coadds = z['COADDS']
        norm = coeff[0]+coeff[1]*fitsarr/coadds+coeff[2]*(fitsarr/coadds)**2
        fitsarr = fitsarr / norm
        z['LINHIST'] = time.asctime()
        if (len(out_file) > 0):
            hl = pyfits.HDUList()
            hl.append(pyfits.ImageHDU(fitsarr,z))
            hl.writeto(out_file,clobber=True)
    return fitsarr

 def _calibration_subarr(self, rdir, flat_file, dark_file, szx, szy, wave=0):
    """A function designed to be used internally only, which chops out the central part
    of calibration data for sub-arrays. It also automatically finds the nearest wavelength
    flat if an appropriate flat doesn't exist. """
    if len(flat_file) > 0:
        try:
            flat = pyfits.getdata(rdir + flat_file,0)
        except:
            if wave>0:
                #Find the flat file with the nearest wavelengths. In this case, ignore
                #corona flats.
                flat_files = glob.glob(rdir + 'flat*fits')
                if len(flat_files)==0:
                    print("No flat! Are you sure that this is your reduction directory? " + rdir)
                    pdb.set_trace()
                waves = []
                for ffile in flat_files:
                    if ffile.find("corona") > 0:
                        waves.append(-1)
                    else:
                        try:
                            wave_file = pyfits.getheader(ffile)['WAVE']
                            waves.append(wave_file)
                        except:
                            print("Missing header keyword WAVE!")
                            pdb.set_trace()
                waves = np.array(waves)
                ix = np.argmin(np.abs(waves - wave))
                new_flat_file = flat_files[ix][len(rdir):]
                print("*** Flat file {0:s} not found! Using {1:s} intead. ***".format(flat_file, new_flat_file))
                flat_file = new_flat_file
                flat = pyfits.getdata(rdir + flat_file,0)
            else:
                print("ERROR - no flat file!")
                pdb.set_trace()
        flat = flat[flat.shape[0]//2 - szy//2:flat.shape[0]//2 + szy//2,flat.shape[1]//2 - szx//2:flat.shape[1]//2 + szx//2]
        bad = pyfits.getdata(rdir + flat_file,1)
        bad = bad[bad.shape[0]//2 - szy//2:bad.shape[0]//2 + szy//2,bad.shape[1]//2 - szx//2:bad.shape[1]//2 + szx//2]
    else:
        flat = np.ones((szy,szx))
        bad = np.zeros((szy,szx))
    if len(dark_file) > 0:
        try:
            dark = pyfits.getdata(rdir + dark_file,0)
            if (szy != dark.shape[0]):
                print("Warning - Dark is of the wrong shape!")
                dark = dark[dark.shape[0]//2 - szy//2:dark.shape[0]//2 + szy//2, \
                       dark.shape[1]//2 - szx//2:dark.shape[1]//2 + szx//2]
        except:
            print("*** Warning - Dark file {0:s} not found! Using zeros for dark ***".format(dark_file))
            dark = np.zeros((szy,szx))
    else:
        dark = np.zeros((szy,szx))
    return (flat,dark,bad)

 def clean_no_dither(self, in_files, fmask_file='',dark_file='', flat_file='', fmask=[],\
     subarr=128, extra_threshold=7, out_file='', median_cut=0.7, destripe=True, ddir='', \
     rdir='', cdir='', manual_click=False):
    """Clean a series of fits files, including: applying the dark, flat,
    removing bad pixels and cosmic rays. This can also be used for dithered data,
    but it will not subtract the dithered positions. There reason for two separate
    programs includes that for dithered data, bad pixel rejection etc has to be done on
    *all* riles.

    Parameters
    ----------
    in_files : array_like (dtype=string).
        A list if input filenames.
    dark_file: string
        The dark file, previously created with make_dark
    flat_file: string
        The flat file, previously created with make_flat
    ftpix: ( (N) array, (N) array)
        The pixels in the data's Fourier Transform that include all non-zero
        values (created using pupil_sampling)
    subarr: int, optional
        The width of the subarray. Defaults to 128 for speed.
    extra_threshold: float, optional
        A threshold for identifying additional bad pixels and cosmic rays.
    outfile: string,optional
        A filename to save the cube as, including the header of the first
        fits file in the cube plus extra information.

    Returns
    -------
    The cube of cleaned frames.
    """
    return self.clean_dithered(in_files, fmask_file=fmask_file,dark_file=dark_file, flat_file=flat_file, fmask=fmask,\
     subarr=subarr,extra_threshold=extra_threshold,out_file=out_file,median_cut=median_cut, destripe=destripe, \
     ddir=ddir, rdir=rdir, cdir=cdir, manual_click=manual_click, dither=False)

 def clean_dithered(self, in_files, fmask_file='',dark_file='', flat_file='', fmask=[],\
    subarr=None, extra_threshold=7,out_file='',median_cut=0.7, destripe=True, \
    manual_click=False, ddir='', rdir='', cdir='', dither=True, show_wait=1, \
    dodgy_badpix_speedup=False, subtract_median=False, extra_diagnostic_plots=False):
    """Clean a series of fits files, including: applying the dark and flat, removing bad pixels and
    cosmic rays, creating a `rough supersky' in order to find a mean image, identifying the target and any
    secondary targets, identifying appropriate sky frames for each frame and subtracting these off. In
    order to find objects in the image in vertical angle mode, the assumption is made that rotation
    is much less than interferogram size.

    To enhance code readability, many of the options in previous routines have been removed.

    Parameters
    ----------
    in_files : array_like (dtype=string).
        A list if input filenames.
    fourier_mask_file: string
        The Fourier mask file - input either this or fmask and subarr
    dark_file: string
        The dark file, previously created with make_dark
    flat_file: string
        The flat file, previously created with make_flat
    ftpix: ( (N) array, (N) array)
        The pixels in the data's Fourier Transform that include all non-zero
        values (created using pupil_sampling)
    subarr: int
        The size of the sub-array, if ftpix is given manually rather than as an fmask_file
    extra_threshold: float, optional
        A threshold for identifying additional bad pixels and cosmic rays.
    outfile: string,optional
        A filename to save the cube as, including the header of the first
        fits file in the cube plus extra information.
    destripe: bool
        Do we destripe the data? This is *bad* for thermal infrared (e.g. the L' filter)
        TODO: Make this a default depending on filter.

    Returns
    -------
    The cube of cleaned frames.
    """
    #Allow over-riding default data, cube and analysis directories.
    if (ddir == ''):
        ddir = self.ddir
    if (rdir == ''):
        rdir = self.rdir
    if (cdir == ''):
        cdir = self.cdir
    #Decide on the image size from the first file. !!! are x and y around the right way?
    try:
        in_fits = pyfits.open(ddir + in_files[0], ignore_missing_end=True)
    except:
        in_fits = pyfits.open(ddir + in_files[0] + '.gz', ignore_missing_end=True)
    h = in_fits[0].header
    in_fits.close()
    szx = h['NAXIS1']
    szy = h['NAXIS2']

    #Extract important information from the header...
    hinfo = self.info_from_header(h, subarr=subarr)
    rnoise = hinfo['rnoise']
    gain = hinfo['gain']
    rad_pixel = hinfo['rad_pixel']
    if (len(dark_file) == 0):
        dark_file = hinfo['dark_file']
    if (len(flat_file) == 0):
        flat_file = hinfo['flat_file']

    #If we set the fmask manually, then don't use the file.
    if len(fmask) == 0:
        #If no file is given, find it automatically.
        if len(fmask_file) == 0:
            fmask_file = hinfo['ftpix_file']
        try:
            fmask = pyfits.getdata(rdir + fmask_file,1)
            subarr = pyfits.getheader(rdir + fmask_file)['SUBARR']
        except:
            print("Error - couldn't find kp/Fourier mask file: " +fmask_file+ " in directory: " + rdir)
            raise UserWarning
    elif not subarr:
        raise UserWarning("Must set subarr if fmask is set!")

    #Allocate memory for the cube and the full cube
    nf = len(in_files)
    cube = np.zeros((nf,subarr,subarr))
    full_cube = np.zeros((nf,szy, szx))
    bad_cube = np.zeros((nf,subarr,subarr), dtype=np.uint8)

    #Chop out the appropriate part of the flat, dark, bad arrays
    (flat,dark,bad) = self._calibration_subarr(rdir, flat_file, dark_file, szx, szy, wave=hinfo['wave'])
    wbad = np.where(bad)
    #Go through the files, cleaning them one at a time and adding to the cube.
    pas = np.zeros(nf)
    raoffs = np.zeros(nf)
    decoffs = np.zeros(nf)
    decs = np.zeros(nf)
    maxs = np.zeros(nf)
    xpeaks = np.zeros(nf,dtype=int)
    ypeaks = np.zeros(nf,dtype=int)
    backgrounds = np.zeros(nf)
    for i in range(nf):
        #First, find the position angles from the header keywords. NB this is the Sky-PA of chip vertical.
        try:
            in_fits = pyfits.open(ddir + in_files[i], ignore_missing_end=True)
        except:
            in_fits = pyfits.open(ddir + in_files[i] + '.gz', ignore_missing_end=True)
        hdr = in_fits[0].header
        in_fits.close()
        pas[i]=360.+hdr['PARANG']+hdr['ROTPPOSN']-hdr['EL']-hdr['INSTANGL']
        raoffs[i]=hdr['RAOFF']
        decoffs[i]=hdr['DECOFF']
        decs[i] =hdr['DEC']

        #Read in the image - making a nonlinearity correction
        im = self.linearize_nirc2(ddir + in_files[i])

        #Find saturated pixels and remove them.
        saturation = self.saturated_pixels(im,hdr)
        bad[saturation]=1
        surrounded = self.is_bad_surrounded(bad)
        bad+=surrounded

        #Destripe, then clean the data using the dark and the flat. This might change
        #the background, so allow for this.
        backgrounds[i] = np.median(im)
        im = self.destripe_nirc2(im, do_destripe=destripe, subtract_median=subtract_median)
        backgrounds[i] -= np.median(im)

        #!!! It is debatable whether the dark on the next line is really useful... but setting
        #dark_file='' removes its effect.
        im = (im - dark)/flat
        #For display purposes, we do a dodgy bad pixel correction.
        mim = nd.filters.median_filter(im,size=3)
        im[bad] = mim[bad]
        full_cube[i,:,:] = im

    #Find the rough "supersky", by taking the 25th percentile of each pixel.
    if (dither):
        rough_supersky = np.percentile(full_cube, 25.0, axis=0)
    else:
        rough_supersky = np.zeros(im.shape)
    #Subtract this supersky off each frame. Don't worry - all strictly pixel-dependent
    #offsets are removed in any case so this doesn't bias the data.
    for i in range(nf):
        full_cube[i,:,:] -= rough_supersky

    #Now the NIRC2-specific stuff...
    raoffs = raoffs - np.mean(raoffs)
    decoffs = decoffs - np.mean(decoffs)
    shifts = np.zeros((nf,2),dtype=int)
    im_mean = np.zeros((szy, szx))
    for i in range(nf):
        th = np.radians(pas[i])
        rot_mat = np.array([[np.cos(th), -np.sin(th)],[-np.sin(th), -np.cos(th)]])
        shifts[i,:] = np.dot(rot_mat, np.radians(np.array([raoffs[i]*np.cos(np.radians(decs[i])), decoffs[i]]))/rad_pixel ).astype(int)
        im_mean = im_mean + np.roll(np.roll(full_cube[i,:,:],-shifts[i,0], axis=1), -shifts[i,1], axis=0)

    #Find the star...
    #Show the image, y-axis reversed.
    plt.clf()
    plt.imshow(np.arcsinh(im_mean/100), interpolation='nearest', origin='lower')
    arrow_xy = np.dot(rot_mat, [0,-30])
    plt.arrow(60,60,arrow_xy[0],arrow_xy[1],width=0.2)
    plt.text(60+1.3*arrow_xy[0], 60+1.3*arrow_xy[1], 'N')
    arrow_xy = np.dot(rot_mat, [-30,0])
    plt.arrow(60,60,arrow_xy[0],arrow_xy[1],width=0.2)
    plt.text(60+1.3*arrow_xy[0], 60+1.3*arrow_xy[1], 'E')
    if manual_click:
        plt.title('Click on target...')
        max_ix = plt.ginput(1, timeout=0)[0]
        #To match the (y,x) order below, change this...
        max_ix = int(max_ix[1]), int(max_ix[0])
    else:
        im_filt = nd.filters.median_filter(im_mean,size=5)
        max_ix = np.unravel_index(im_filt.argmax(), im_filt.shape)
        plt.title('Identified target shown')
    plt.plot(max_ix[1], max_ix[0], 'wx', markersize=20,markeredgewidth=2)
    plt.pause(0.001)
    #plt.draw()
    print("Maximum x,y: " + str(max_ix[1])+', '+ str(max_ix[0]))
    time.sleep(show_wait)

    #Set the xpeaks and ypeaks values (needed for sub-arraying later)
    for i in range(nf):
        xpeaks[i] = max_ix[1] + shifts[i,0]
        ypeaks[i] = max_ix[0] + shifts[i,1]

    subims = np.empty( (nf,subarr,subarr) )
    #Sky subtract and fix bad pixels.
    for i in range(nf):
        #For each frame, cut out a sub-image centered on this particular (x,y) location
        #corresponding to frame i.
        for j in range(nf):
            #Undo the flat, to minimise the effects of errors in the flat.
            im = full_cube[j,:,:]*flat
            #Roll all the sub-images, and cut them out.
            subims[j,:,:] = np.roll(np.roll(im,subarr//2-ypeaks[i],axis=0),
                subarr//2-xpeaks[i],axis=1)[0:subarr,0:subarr]
        #Find a flat to re-apply
        subflat = np.roll(np.roll(flat,subarr//2-ypeaks[i],axis=0),subarr//2-xpeaks[i],axis=1)
        subflat = subflat[0:subarr,0:subarr]
        #Find the frames that are appropriate for a dither...
        if (dither):
            w = np.where( (xpeaks - xpeaks[i])**2 + (ypeaks - ypeaks[i])**2 > 0.5*subarr**2 )[0]
            if len(w) == 0:
                print("Error: Can not find sky from dithered data - use dither=False for {0:s}".format(in_files[0]))
                return
            #To avoid too many extra bad pixels, we'll use a median here.
            sky = np.median(subims[w,:,:], axis=0)
            backgrounds[i] += np.median(sky)
            #Subtract the sky then re-apply the flat.
            subim = (subims[i,:,:] - sky)/subflat
        else:
            subim = subims[i,:,:]/subflat
        #Find the peak from the sub-image.
        im_filt = nd.filters.median_filter(subim,size=5)
        max_ix = np.unravel_index(im_filt.argmax(), im_filt.shape)
        maxs[i] = subim[max_ix[0],max_ix[1]]

        #subbad is the set of bad pixels in the sub-array.
        subbad = np.roll(np.roll(bad,subarr//2-ypeaks[i],axis=0),subarr//2-xpeaks[i],axis=1)
        subbad = subbad[0:subarr,0:subarr]
        new_bad = subbad.copy()
        subim[np.where(subbad)] = 0
        plt.clf()
        plt.imshow(np.maximum(subim,0)**0.5,interpolation='nearest')
        plt.title(hinfo['targname'])
        plt.pause(0.001)
        #import pdb; pdb.set_trace()
        #plt.draw()
        #Iteratively fix the bad pixels and look for more bad pixels...
        for ntry in range(1,15):
            # Correct the known bad pixels. Ideally, we self-consistently correct
            # all bad pixels at once.
            if dodgy_badpix_speedup:
                self.fix_bad_pixels(subim,new_bad,fmask)
            else:
                self.fix_bad_pixels(subim,subbad,fmask)

            # Search for more bad pixels. Lets use a Fourier technique here, where we
            # take the inverse Fourier transform of the region of the image Fourier transform
            # that is the null space of the MTF
            extra_bad_ft = np.fft.rfft2(subim)*fmask
            bad_image = np.real(np.fft.irfft2(extra_bad_ft))
            mim = nd.filters.median_filter(subim,size=5)

            # NB The next line *should* take experimentally determined readout noise into account
            # rather than a fixed readout noise!!!
            total_noise = np.sqrt(np.maximum((backgrounds[i] + mim)/gain + rnoise**2,rnoise**2))
            bad_image = bad_image/total_noise

            #In case of a single bad pixel, we end up with a ringing effect where the
            #surrounding pixels also look bad. So subtract a median filtered image.
            unsharp_masked = bad_image-nd.filters.median_filter(bad_image,size=3)

            # The "extra_threshold" value for extra bad pixels is a scaling of the median
            # absolute deviation. We set a limit where new bad pixels can't have
            # absolute values morer than 0.2 times the peak bad pixel.
            current_threshold = np.max([0.25*np.max(np.abs(unsharp_masked[new_bad == 0])), \
                extra_threshold*np.median(np.abs(bad_image))])
            extra_bad = np.abs(unsharp_masked) > current_threshold
            n_extra_bad = np.sum(extra_bad)
            print(str(n_extra_bad)+" extra bad pixels or cosmic rays identified. Attempt: "+str(ntry))

            #TESTING - too many bad pixels are identified in L' data, and noise isn't
            #consistent.
            if extra_diagnostic_plots:
                plt.clf()
                plt.imshow(np.maximum(subim,0)**0.5, interpolation='nearest', cmap=cm.cubehelix)
                new_bad_yx = np.where(new_bad)
                plt.plot(new_bad_yx[1], new_bad_yx[0], 'wx')
                plt.axis([0,128,0,128])
                import pdb; pdb.set_trace()

            subbad += extra_bad
            if (ntry == 1):
                new_bad = extra_bad
            else:
                new_bad += extra_bad
                new_bad = extra_bad>0
            if (n_extra_bad == 0):
                break
        print(str(np.sum(subbad)) + " total bad pixels.")

        #Now re-correct both the known and new bad pixels at once.
        self.fix_bad_pixels(subim,subbad,fmask)
        plt.imshow(np.maximum(subim,0)**0.5,interpolation='nearest')
        plt.pause(0.001)
        #plt.draw()

        #Save the data and move on!
        cube[i]=subim
        subbad = subbad>0
        bad_cube[i]=subbad.astype(np.uint8)

    #Find bad frames based on low peak count.
    good = np.where(maxs > median_cut*np.median(maxs))
    good = good[0]
    if (len(good) < nf):
        print(nf-len(good), "  frames rejected due to low peak counts.")
    cube = cube[good,:,:]
    nf = np.shape(cube)[0]
    #If a filename is given, save the file.
    if (len(out_file) > 0):
        hl = pyfits.HDUList()
        h['RNOISE'] = rnoise
        h['PGAIN'] = gain #P means python
#        h['PFILTER'] = !!! Not complete !!!
        h['SZX'] = szx
        h['SZY'] = szy
        h['DDIR'] = ddir
        h['TARGNAME'] = hinfo['targname']
        #NB 'TARGNAME' is the standard target name.
        for i in range(nf):
            h['HISTORY'] = 'Input: ' + in_files[i]
        hl.append(pyfits.ImageHDU(cube,h))
        #Add in the original peak pixel values, forming the image centers in the cube.
        #(this is needed for e.g. undistortion)
        col1 = pyfits.Column(name='xpeak', format='E', array=xpeaks)
        col2 = pyfits.Column(name='ypeak', format='E', array=ypeaks)
        col3 = pyfits.Column(name='pa', format='E', array=pas)
        col4 = pyfits.Column(name='max', format='E', array=maxs)
        col5 = pyfits.Column(name='background', format='E', array=backgrounds)
        cols = pyfits.ColDefs([col1, col2,col3,col4,col5])
        hl.append(pyfits.BinTableHDU.from_columns(cols))
        hl.append(pyfits.ImageHDU(bad_cube))
        hl.writeto(cdir+out_file,clobber=True)
    return cube

if(0):
    n2 = NIRC2()

#Testing destripe only.
if (0):
    f = pyfits.open(file, ignore_missing_end=True)
    im = f[0].data
    f.close()
    plt.imshow(np.minimum(np.maximum(n2.destripe_nirc2(im),-50),50), interpolation='nearest', cmap=cm.gray)
    f.close()

#Testing darks and flats.
if (0):
    dir =  '/Users/mireland/data/nirc2/131115/n'
    extn = '.fits.gz'
    files = [(dir + '{0:04d}' + extn).format(i) for i in range(40,45)]
    #files.extend([(dir + '{0:04d}' + extn).format(i) for i in range(56,61)])
    n2.make_dark(files,'dark.fits')
    files = [(dir + '{0:04d}' + extn).format(i) for i in range(29,40)]
    n2.make_flat(files,'dark.fits','flat.fits')

import numpy,import scipy,import astropy,from astropy


import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import numpy as np
from src.utilities.yolov5.experimental import attempt_load
from src.utilities.yolov5.datasets import LoadStreams, LoadImages
from src.utilities.yolov5.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    scale_coords, xyxy2xywh, strip_optimizer
from src.utilities.yolov5.plots import plot_one_box
from src.utilities.yolov5.torch_utils import select_device, load_classifier, time_synchronized

import config

def load_model(weights=config.WEIGHT_PATH ,device=config.DEVICE):
    select_d = select_device(device)
    return attempt_load([weights], map_location=select_d)

model = load_model()

def detect( image_path ,model=model, image_size = config.IMAGE_SIZE,s_device=config.DEVICE):

    # Load model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(image_size, s=stride)  # check img_size

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    device = select_device(s_device)
    # Run inference
    # if device == 'cpu':
    #     model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once

    t0 = time.time()
    dataset = LoadImages(image_path, img_size=imgsz, stride=stride)

    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img/255
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img, augment=None)[0]

        # Apply NMS
        pred = non_max_suppression(pred, config.CONF_THRESHOLD, config.IOU_THRESHOLD)
        t2 = time_synchronized()

        output = []
        # Process detections
        for i, det in enumerate(pred):  # detections per image
            p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
            p = Path(p)  # to Path

            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                for *xyxy, conf, cls in reversed(det):
                    #label = f'{names[int(cls)]} {conf:.2f}'
                    bbox = {'vertices':[]}

                    bbox['vertices'].append({'x': int(xyxy[0]), 'y': int(xyxy[1])})
                    bbox['vertices'].append({'x': int(xyxy[2]), 'y': int(xyxy[1])})
                    bbox['vertices'].append({'x': int(xyxy[2]), 'y': int(xyxy[3])})
                    bbox['vertices'].append({'x': int(xyxy[0]), 'y': int(xyxy[3])})

                    output.append({'class':names[int(cls)], 'boundingBox':bbox  ,'conf' :round(float(conf),2) })
                    #plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

            print(f'{s}Done. ({t2 - t1:.3f}s)')

            # Stream results
            # view_img =True
            # save_txt =True
            # for *xyxy, conf, cls in reversed(det):
            #     if save_txt:  # Write to file
            #         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
            #         # line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
            #         # with open(txt_path + '.txt', 'a') as f:
            #         #     f.write(('%g ' * len(line)).rstrip() % line + '\n')
            #
            #     if view_img:  # Add bbox to image
            #         label = f'{names[int(cls)]} {conf:.2f}'
            #         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
            #
            # if view_img:
            #     cv2.imwrite('/home/dhiraj/b.png',im0)
            #     cv2.imshow(str(p), im0)
            #     cv2.waitKey(1)  # 1 millisecond

            #Save results (image with detections)

    print(f'Done. ({time.time() - t0:.3f}s)')
    return output


import os
import logging
logger = logging.getLogger(__name__)

import numpy as np
import astropy.io.fits as fits
import matplotlib.pyplot as plt

from ...echelle.imageproc import combine_images
from ...echelle.trace import find_apertures, load_aperture_set
from .common import (get_bias, get_mask, correct_overscan,
                    fix_badpixels,
                    TraceFigure, BackgroundFigure,
                    )

def reduce_feros(config, logtable):
    """Reduce the single fiber data of FEROS.

    Args:
        config (:class:`configparser.ConfigParser`): Config object.
        logtable (:class:`astropy.table.Table`): Table of observing log.

    """
    # extract keywords from config file
    section      = config['data']
    rawpath      = section.get('rawpath')
    statime_key  = section.get('statime_key')
    exptime_key  = section.get('exptime_key')
    direction    = section.get('direction')

    section     = config['reduce']
    midpath     = section.get('midpath')
    odspath     = section.get('odspath')
    figpath     = section.get('figpath')
    mode        = section.get('mode')
    fig_format  = section.get('fig_format')
    oned_suffix = section.get('oned_suffix')
    ncores      = section.get('ncores')

    # create folders if not exist
    if not os.path.exists(figpath): os.mkdir(figpath)
    if not os.path.exists(odspath): os.mkdir(odspath)
    if not os.path.exists(midpath): os.mkdir(midpath)

    # determine number of cores to be used
    if ncores == 'max':
        ncores = os.cpu_count()
    else:
        ncores = min(os.cpu_count(), int(ncores))

    ################ parse bias ########################
    bias, bias_card_lst = get_bias(config, logtable)

    ############### find flat groups ##################

    # initialize flat_groups for single fiber
    flat_groups = {}
    # flat_groups = {'flat_M': [fileid1, fileid2, ...],
    #                'flat_N': [fileid1, fileid2, ...]}

    for logitem in logtable:
        if logitem['object']=='FLAT' and logitem['binning']=='(1, 1)':

            # find a proper name for this flat
            flatname = '{:g}'.format(logitem['exptime'])

            # add flatname to flat_groups
            if flatname not in flat_groups:
                flat_groups[flatname] = []
            flat_groups[flatname].append(logitem)

    ################# Combine the flats and trace the orders ###################

    # first combine the flats
    for flatname, logitem_lst in flat_groups.items():
        nflat = len(logitem_lst)       # number of flat fieldings

        flat_filename    = os.path.join(midpath,
                            'flat_{}.fits'.format(flatname))
        aperset_filename = os.path.join(midpath,
                            'trace_flat_{}.trc'.format(flatname))
        aperset_regname  = os.path.join(midpath,
                            'trace_flat_{}.reg'.format(flatname))
        trace_figname = os.path.join(figpath,
                        'trace_flat_{}.{}'.format(flatname, fig_format))

        # get flat_data and mask_array
        if mode=='debug' and os.path.exists(flat_filename) \
            and os.path.exists(aperset_filename):
            pass
        else:
            # if the above conditions are not satisfied, comine each flat
            data_lst = []
            head_lst = []
            exptime_lst = []

            print('* Combine {} Flat Images: {}'.format(nflat, flat_filename))
            fmt_str = '  - {:>7s} {:^23} {:^8s} {:^7} {:^8} {:^6}'
            head_str = fmt_str.format('frameid', 'FileID', 'Object', 'exptime',
                        'N(sat)', 'Q95')

            for iframe, logitem in enumerate(logitem_lst):
                # read each individual flat frame
                fname = 'FEROS.{}.fits'.format(logitem['fileid'])
                filename = os.path.join(rawpath, fname)
                data, head = fits.getdata(filename, header=True)
                exptime_lst.append(head[exptime_key])
                mask = get_mask(data, head)
                sat_mask = (mask&4>0)
                bad_mask = (mask&2>0)
                if iframe == 0:
                    allmask = np.zeros_like(mask, dtype=np.int16)
                allmask += sat_mask

                # correct overscan for flat
                data, card_lst = correct_overscan(data, head)
                for key, value in card_lst:
                    head.append((key, value))

                # correct bias for flat, if has bias
                if bias is None:
                    message = 'No bias. skipped bias correction'
                else:
                    data = data - bias
                    message = 'Bias corrected'
                logger.info(message)

                # print info
                if iframe == 0:
                    print(head_str)
                message = fmt_str.format(
                            '[{:d}]'.format(logitem['frameid']),
                            logitem['fileid'], logitem['object'],
                            logitem['exptime'],
                            logitem['nsat'], logitem['q95'])
                print(message)

                data_lst.append(data)

            if nflat == 1:
                flat_data = data_lst[0]
            else:
                data_lst = np.array(data_lst)
                flat_data = combine_images(data_lst,
                                mode       = 'mean',
                                upper_clip = 10,
                                maxiter    = 5,
                                maskmode   = (None, 'max')[nflat>3],
                                ncores     = ncores,
                                )
            fig = plt.figure(dpi=300)
            ax = fig.gca()
            ax.plot(flat_data[2166, 0:400],lw=0.5, color='C0')
            # fix badpixels in flat
            flat_data = fix_badpixels(flat_data, bad_mask)
            ax.plot(flat_data[2166, 0:400],lw=0.5, color='C1')
            plt.show()

            # get mean exposure time and write it to header
            head = fits.Header()
            exptime = np.array(exptime_lst).mean()
            head[exptime_key] = exptime

            # find saturation mask
            sat_mask = allmask > nflat/2.
            flat_mask = np.int16(sat_mask)*4 + np.int16(bad_mask)*2

            # get exposure time normalized flats
            flat_norm = flat_data/exptime

            # create the trace figure
            tracefig = TraceFigure()

            section = config['reduce.trace']
            aperset = find_apertures(flat_data, flat_mask,
                        transpose  = True,
                        scan_step  = section.getint('scan_step'),
                        minimum    = section.getfloat('minimum'),
                        separation = section.get('separation'),
                        align_deg  = section.getint('align_deg'),
                        filling    = section.getfloat('filling'),
                        degree     = section.getint('degree'),
                        conv_core  = 20,
                        display    = section.getboolean('display'),
                        fig        = tracefig,
                        )

            # save the trace figure
            tracefig.adjust_positions()
            title = 'Trace for {}'.format(flat_filename)
            tracefig.suptitle(title, fontsize=15)
            tracefig.savefig(trace_figname)

            aperset.save_txt(aperset_filename)
            aperset.save_reg(aperset_regname)

            # do the flat fielding
            # prepare the output mid-prococess figures in debug mode
            if mode=='debug':
                figname = 'flat_aperpar_{}_%03d.{}'.format(
                            flatname, fig_format)
                fig_aperpar = os.path.join(figpath, figname)
            else:
                fig_aperpar = None

            # prepare the name for slit figure
            figname = 'slit_flat_{}.{}'.format(flatname, fig_format)
            fig_slit = os.path.join(figpath, figname)

            # prepare the name for slit file
            fname = 'slit_flat_{}.dat'.format(flatname)
            slit_file = os.path.join(midpath, fname)

            #section = config['reduce.flat']

            # pack results and save to fits
            hdu_lst = fits.HDUList([
                        fits.PrimaryHDU(flat_data, head),
                        fits.ImageHDU(flat_mask),
                        fits.ImageHDU(flat_norm),
                        #fits.ImageHDU(flat_sens),
                        #fits.BinTableHDU(flat_spec),
                        ])
            hdu_lst.writeto(flat_filename, overwrite=True)

            # now flt_data and mask_array are prepared


import random
import json, sys
import torch
import argparse
import numpy as np
import yaml
import os

sys.path.append("./")
from model import GenreClassficationModel
from utils import bag_of_words, preproces_text

with open("src/config.yaml", "r") as file:
    config = yaml.safe_load(file)

parser = argparse.ArgumentParser(description="movie title and description")
parser.add_argument("--title", type=str, required=True, help="the movie title")
parser.add_argument("--description", type=str, required=True, help="description of the movie")
args = parser.parse_args()
MovieDescription = args.description
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
FILE = config["File"]
model = torch.load(FILE)
# Retrieve the model's hyperparameters
input_size = model["input_size"]
hidden_size = model["hidden_size"]
output_size = model["output_size"]
all_words = model["all_words"]
model_state = model["model_state"]
classes = model["classes"]
# Load trained model's parameters
model = GenreClassficationModel(input_size, hidden_size, output_size).to(device)
model.load_state_dict(model_state)
model.eval()
# We generate the prediction
X = preproces_text(MovieDescription)
X = bag_of_words(X, all_words)
X = X.reshape(1, X.shape[0])
X = torch.from_numpy(X).to(device)
output = model(X)
_, predicted = torch.max(output, dim=1)
probs = torch.softmax(output, dim=1)
response = {
    "title": args.title,
    "description": MovieDescription,
    "genre": classes[predicted.tolist()[0]],
}
print(response)