# TradingPairsStrategyExample.py

# This is an example program for a class in order
# to learn how to write code at quantopian.com.
# DO NOT USE THIS CODE FOR REAL TRADING!  IT IS NOT
# A GOOD STRATEGY!

# This program trades correlated pairs (stocks that
# general move together to a high degree) when their
# correlation starts to drift.  The program assumes
# that the correlation will correct itself, and that
# the 2 stocks will go back to their matching movements.
# We their movement is out of sync, we short the over performing
# one, and buy the under performing one on the assumption
# that they will get back to lock step for their movement.

# One of the reasons, that we are choosing 2 stocks, is
# that we are trying to find how the "industry" is performing
# as a whole.  I guess you can use > 2 stocks to get a better
# picture of the airlines industry.  M/b a airlines index.

import numpy

def initialize(context):
    # Run when the program starts up.

    # Stocks to trade:
    context.aa = sid(45971)    # American Airlines
    context.ual = sid(28051)   # United Airlines

    context.long_on_spread = False
    context.short_on_spread = False

    # For this example, we are arbitralily running this
    # 60 minutes b/4 market closed.  In real life, I would
    # have picked 30-60 minutes after market open.
    schedule_function(check_pairs, date_rules.every_day(),
                      time_rules.market_close(minutes=60))

def check_pairs(context, data):

    # This function will be scheduled.  It finds when the correlation
    # has drifted, and shorts the overperformer, and buys the underperformer.
    # When the correlation gets back in normal, the function closes
    # the positions.

    # Looks like we are hard coding which stocks we are using into this
    # function.  I am not a fan of this.
    aa = context.aa
    ual = context.ual

    # Get the last 30 days of closing price info for our 2 stocks:
    # "1d" = 1 day, aka daily data.
    prices_DF = data.history([aa, ual], "price", 30, "1d")

    # Not sure why we need this.  I guess I'll learn more later.
    yesterdays_close_prices_list = prices_DF.iloc[-1:]

    # Get a list of the spread for the last 30 days,
    # and then get the average value of that list.
    # So, this is 1 data point.
    Spread_Ave_30 = numpy.mean(prices_DF[aa] - prices_DF[ual])

    # Get a list of the spread for the last 30 days,
    # and find the standard deviation for that list.
    # So, this is 1 data point.
    Std_Dev_Ave_30 = numpy.std(prices_DF[aa] - prices_DF[ual])

    print(Std_Dev_Ave_30)

    # Average of the 2 prices from yesterday's close:
    ave_of_closing_prices = numpy.mean(yesterdays_close_prices_list[aa] - yesterdays_close_prices_list[ual])

    # Normalize the deviation.
    # Find the difference b/t yesterday's spread and the average spread for the past 30 days.
    # And, then normalize it by divided by the std. dev. based off of the last 30 days.
    z_score = (ave_of_closing_prices - Spread_Ave_30)/Std_Dev_Ave_30

    # The z_score is a measure of how well the stocks are moving together.
    # Close to 0 means in sync.

    # If we have positions open, and our z_score indicates that the
    # stocks are moving back together, then close the open orders b/c
    # our advantage is now gone for the time being.
    if ( abs(z_score) < 0.1 ) and ( context.long_on_spread or context.short_on_spread ):
        # Close the orders if any are open:
        order_target_percent(aa, 0)
        order_target_percent(ual, 0)

        # Let the program know we don't have any open positions:
        context.long_on_spread = False
        context.short_on_spread = False

    # spread = AA - UAL
    # So, if z_score gets excessively positive, then AA is outperforming its average.
    # So, short it on the assumption it will return to its normal value.
    # Short AA: b/c it's outperforming
    # Buy UAL: b/c it may be underperforming, or will just catch up to AA.
    # Open a trade if we are not already in one:

    # I would do this differently.  I would think that you would, first, have a
    # test to close open positions if you are currently long, and the
    # z_score goes back to < 1.0.

    if z_score > 1.0 and not context.short_on_spread:

        # 50% short AA.  50% buy United.
        order_target_percent(aa, -0.5)
        order_target_percent(ual, 0.5)

        # Make a note that we shorted AA:
        context.short_on_spread = True
        context.long_on_spread = False

    # Are conditions right to go long w/ AA, and short United
    # Is United outperforming its average?
    if z_score < -1.0 and not context.long_on_spread:
        order_target_percent(aa, 0.5)
        order_target_percent(ual, -0.5)

        # Make a note that we bought AA:
        context.short_on_spread = False
        context.long_on_spread = True

    # Log the z score:
    record(z_score = z_score)


""" use> python csv2gz.py filename
for training set """

import theano
import pickle
import numpy
import gzip
import sys, time

def csv2gz(filename):
    t1 = time.time()
    print('='*80)
    print('The', filename, 'will be converted to', filename +'.pkl.gz', '\n')

    my_train_x = []
    my_train_y = []
    my_valid_x = []
    my_valid_y = []
    my_test_x = []
    my_test_y = []

    # check the size of csv.file
    data = numpy.genfromtxt(filename, delimiter=',')
    numberD, numberinput = data.shape
    print('The training set number and size of input = ', data.shape)
    print('This set will be divided into three parts as training set (0.8), validation set(0.1) and test est(0.1)')
    print('The size of training set = ', 0.8*numberD)
    print('The size of validation set = ', 0.1*numberD)
    print('The size of test set = ', 0.1*numberD, '\n')

    # for i, line in enumerate(fp):
    with open(filename, 'r') as f:
        for i, line in enumerate(f):
            if i == 0:
                print('The first one size = ', len(line.split(',')))
            if i < 0.8*numberD:
                my_list = line.split(',') # replace with your own separator instead
                my_train_x.append(my_list[1:]) # omitting identifier in [0] and target in [-1]
                my_train_y.append(my_list[0])
            elif 0.8*numberD <= i < 0.9*numberD:
                my_list = line.split(',')  # replace with your own separator instead
                my_valid_x.append(my_list[1:])  # omitting identifier in [0] and target in [-1]
                my_valid_y.append(my_list[0])
            else:
                my_list = line.split(',')  # replace with your own separator instead
                my_test_x.append(my_list[1:])  # omitting identifier in [0] and target in [-1]
                my_test_y.append(my_list[0])

    # train_set, valid_set, test_set format: tuple(input, target)
    # input is a numpy.ndarray of 2 dimensions (a matrix)
    # where each row corresponds to an example. target is a
    # numpy.ndarray of 1 dimension (vector) that has the same length as
    # the number of rows in the input. It should give the target
    # to the example with the same index in the input.

    print('training set = ', len(my_train_x), len(my_train_x[0]), len(my_train_y), len(my_train_y[0]))
    print('valid set = ', len(my_valid_x), len(my_valid_x[0]), len(my_valid_y), len(my_valid_y[0]))
    print('test set = ', len(my_test_x), len(my_test_x[0]), len(my_test_y), len(my_test_y[0]))

    f = open(filename+'.pkl', 'wb')
    pickle.dump([(my_train_x, my_train_y), (my_valid_x, my_valid_y), (my_test_x, my_test_y)], f)
    f.close()

    f_in = open(filename+'.pkl', 'rb')
    f_out = gzip.open(filename+'.pkl.gz', 'wb')
    f_out.writelines(f_in)
    f_out.close()
    f_in.close()

    t2 = time.time()
    print('time = ', t2 - t1, ' s ')
    print('time = ', (t2 - t1) / 60, ' mins ')
    print(time.asctime(time.localtime(time.time())))
    print('='*80)


import cv2
import numpy as np
from bs import ScreenGrabber, Vector2
from ..typing import MazeLoader, Maze, MazeFactory

def fill_in_missing_lines(lines):
    deltas = []
    last_v = 0
    for v in lines:
        deltas.append(v - last_v)
        last_v = v

    _lines = []
    tg = sorted(deltas)[len(deltas)//2]

    f = 0
    for i, (d, v) in enumerate(zip(deltas, lines)):
        if d - tg < tg:
            _lines.append(v)
            continue
        for _ in range(1, round(d/tg)):
            _lines.append(_lines[-1] + tg)
        _lines.append(v)
    return _lines, tg

def find_vlines(img):

    imgl = cv2.filter2D(img, -1, np.array([
        [1, 0, -1],
        [2, 0, -2],
        [1, 0, -1],
    ]))
    imgr = cv2.filter2D(img, -1, np.array([
        [-1, 0, 1],
        [-2, 0, 2],
        [-1, 0, 1],
    ]))

    vlines = [0] * img.shape[1]
    for y, row in enumerate(img):
        for x, _ in enumerate(row):
            try:
                if all([
                    imgl[i+y, x] == 255
                    for i in range(-2, 2)
                ]) or all([
                    imgr[i+y, x] == 255
                    for i in range(-2, 2)
                ]):
                    vlines[x] += 1
            except:
                pass

    ret = []
    p = 0
    cooldown = 0
    last_tw = 0
    for x, v in enumerate(vlines):
        cooldown = max(0, cooldown - 1)
        if v > p+20 and cooldown == 0:
            ret.append(x)
            cooldown = last_tw // 3
            last_tw = 0
        else:
            last_tw += 1
        p = v

    #if len(ret) % 2 != 0:
        #ret = ret[:-1]

    return fill_in_missing_lines(ret)

def find_hlines(img):

    imgt = cv2.filter2D(img, -1, np.array([
        [1, 2, 1],
        [0, 0, 0],
        [-1, -2, -1],
    ]))
    imgb = cv2.filter2D(img, -1, np.array([
        [-1, -2, -1],
        [0, 0, 0],
        [1, 2, 1],
    ]))

    lines = [0] * img.shape[0]
    for y, row in enumerate(img):
        for x, _ in enumerate(row):
            try:
                if all([
                    imgt[y, i+x] > 200
                    for i in range(-2, 2)
                ]) or all([
                    imgb[y, i+x] > 200
                    for i in range(-2, 2)
                ]):
                    lines[y] += 1
            except:
                pass

    ret = []
    p = 0
    cooldown = 0
    last_tw = 0
    for y, v in enumerate(lines):
        cooldown = max(0, cooldown - 1)
        if v > p+10 and cooldown == 0:
            ret.append(y)
            cooldown = last_tw // 2
            last_tw = 0
        else:
            last_tw += 1
        p = v

    if len(ret) % 2 != 0:
        ret = ret[:-1]

    return fill_in_missing_lines(ret)

class EdgeDetectingMazeLoader(MazeLoader):
    """Load a maze from a screen shot and use edge-detection to find build the maze"""

    def __init__(self, screen: ScreenGrabber, factory: MazeFactory):
        self._screen = screen
        self._factory = factory

    def load(self) -> Maze:
        img = self._screen.grab_gray()

        # Find the map in the screen capture
        th, tw = img.shape
        hslice = [i for i in range(tw) if img[th//2, i] == 189]
        vslice = [i for i in range(th) if img[i, tw//2] == 189]
        xs, xe = hslice[0], hslice[-1]
        ys, ye = vslice[0], vslice[-1]
        img = img[ys:ye,xs:xe]

        # Make the image smaller to speed things up
        th, tw = [int(x*0.5) for x in img.shape]
        img = cv2.resize(img, (tw, th))

        # Find all the vertical and horizontal lines.
        vlines, ts = find_vlines(img)
        hlines, _ = find_hlines(img)

        # Push in a leading edge to simply the math.
        vlines.insert(0, 0)
        hlines.insert(0, 0)

        # Compute the step-width as half the tile size.
        ts //= 2

        #for x in vlines:
        #    img[:, x] = 0
        #for y in hlines:
        #    img[y, :] = 0

        #while True:
        #    cv2.imshow(f"img", img)
        #    if cv2.waitKey(25) & 0xFF == ord("q"):
        #        cv2.destroyAllWindows()
        #        break

        # Fill a grid by checking the center pixel of each tile.
        grid = np.zeros((len(hlines), len(vlines)), dtype="uint8")
        for gy, y in enumerate(hlines):
            for gx, x in enumerate(vlines):
                if img[y+ts, x+ts] == 189:
                    grid[gy, gx] = 1

        th, tw = grid.shape
        return self._factory.create(grid, Vector2(1, 1), Vector2(tw-2, th-2))


import wx,sys,subprocess,os
from BabelBrushVolumeHub import BabelBrushVolumeHub,create_directory_structure

class VolumeList(wx.ListCtrl):
    def __init__(self, parent,volume_hub):
        wx.ListCtrl.__init__(
            self, parent, -1,
            style=wx.LC_REPORT|wx.LC_VIRTUAL|wx.LC_HRULES|wx.LC_VRULES
            )

        self.volume_hub=volume_hub
        self.il = wx.ImageList(60,64)

        for item in volume_hub.volumes:
            item["_image_index"]=self.il.Add(wx.Bitmap(item["thumbnail"],wx.BITMAP_TYPE_JPEG))

        self.SetImageList(self.il, wx.IMAGE_LIST_SMALL)

        self.InsertColumn(0, "Thumb")
        self.InsertColumn(1, "Volume")
        self.InsertColumn(2, "File")

        self.SetColumnWidth(0, 65)
        self.SetColumnWidth(1, 100)
        self.SetColumnWidth(2, 180)

        self.SetItemCount(len(volume_hub.volumes))

    def item_added(self):

        last_added= self.volume_hub.volumes[-1]
        last_added["_image_index"]=self.il.Add(wx.Bitmap(last_added["thumbnail"],wx.BITMAP_TYPE_JPEG))
        self.SetItemCount(len(self.volume_hub.volumes))

    def hub_changed(self):
        self.il.RemoveAll()
        for item in self.volume_hub.volumes:
            item["_image_index"]=self.il.Add(wx.Bitmap(item["thumbnail"],wx.BITMAP_TYPE_JPEG))
        self.SetItemCount(len(self.volume_hub.volumes))
        self.Refresh()

    def OnGetItemText(self, item, col):
        if col==1:
            return self.volume_hub.volumes[item]["_stub"]
        elif col==2:
            return self.volume_hub.volumes[item]["Image Name"]
        return ""

    def OnGetItemImage(self, item):
        if item >=  len(self.volume_hub.volumes):
            return None
        return self.volume_hub.volumes[item]["_image_index"]

class BabelBrushVolumeManager(wx.Frame):

    def __init__(self,base_folder=None):
        parent = None
        title = "BabelBrush Manager"
        super(BabelBrushVolumeManager, self).__init__(parent, title = title,size=(1000,600))

        label_font=  wx.Font(16, wx.DECORATIVE, wx.NORMAL, wx.NORMAL)
        info_font= wx.Font(13,wx.FONTFAMILY_SWISS,wx.NORMAL, wx.NORMAL)
        norm_font= wx.Font(12,wx.FONTFAMILY_SWISS,wx.NORMAL, wx.NORMAL)
        norm_font_bold= wx.Font(12,wx.FONTFAMILY_SWISS,wx.NORMAL, wx.BOLD)
        #splt top bottom
        outer_box=wx.BoxSizer(wx.VERTICAL)
        top_box=wx.BoxSizer(wx.HORIZONTAL)
        outer_box.Add(top_box,0,wx.ALL|wx.EXPAND,10)

        #top bar
        folder_label=wx.StaticText(self,label="Hub Folder:")
        folder_label.SetFont(label_font)
        top_box.Add(folder_label,0,wx.RIGHT,15)

        self.hub_label = wx.StaticText(self,size=(300,40),style=wx.ST_ELLIPSIZE_END)
        top_box.Add(self.hub_label)

        change_folder_button = wx.Button(self,-1,"Open")
        change_folder_button.Bind(wx.EVT_BUTTON,self.change_hub)
        top_box.Add(change_folder_button,0,wx.ALL,3)

        new_folder_button = wx.Button(self,-1,"Create New")
        new_folder_button.Bind(wx.EVT_BUTTON,self.create_new_hub)
        top_box.Add(new_folder_button,0,wx.ALL,3)

        self.launch_button = wx.Button(self,-1,"Launch Babel Brush")
        self.launch_button.Bind(wx.EVT_BUTTON,self.launch_babel)
        top_box.Add(self.launch_button,0,wx.ALL,3)

        main_box= wx.BoxSizer(wx.HORIZONTAL)
        outer_box.Add(main_box,1,wx.EXPAND)

        #spli left right
        left_box= wx.BoxSizer(wx.VERTICAL)
        main_box.Add(left_box,1,wx.EXPAND|wx.ALL,5)
        right_box = wx.BoxSizer(wx.VERTICAL)
        main_box.Add(right_box,1,wx.EXPAND|wx.ALL,5)

        list_box_menu= wx.BoxSizer(wx.HORIZONTAL)
        list_label=wx.StaticText(self,label="Volumes")
        list_label.SetFont(label_font)
        list_box_menu.Add(list_label,1,wx.ALL,5)

        self.index_button = wx.Button(self,-1,"Add nii file")
        self.index_button.Bind(wx.EVT_BUTTON,self.addIndex)

        self.del_button= wx.Button(self,-1,"Delete")
        self.del_button.Bind(wx.EVT_BUTTON,self.delete_volume)
        list_box_menu.Add(self.index_button,0,wx.ALL,5)
        list_box_menu.Add(self.del_button,0,wx.ALL,5)
        right_box.Add(list_box_menu,0)

        main_left_box= wx.BoxSizer(wx.VERTICAL)
        bottom_left_box=wx.BoxSizer(wx.VERTICAL)
        left_box.Add(main_left_box,4,wx.EXPAND|wx.ALL)
        left_box.Add(bottom_left_box,1,wx.EXPAND|wx.ALL)

        config_box= wx.BoxSizer(wx.HORIZONTAL)
        fiji_lab= wx.StaticText(self,label="Fiji Path:")
        fiji_lab.SetFont(norm_font_bold)
        config_box.Add(fiji_lab)
        self.fiji_label=wx.StaticText(self,label="None",size=(250,-1),style=wx.ST_ELLIPSIZE_END)
        self.fiji_label.SetFont(norm_font)
        config_box.Add(self.fiji_label)
        self.fiji_btn=wx.Button(self,-1,"Change")
        self.fiji_btn.Bind(wx.EVT_BUTTON,self.change_fiji_path)
        self.fiji_btn.Disable()
        config_box.Add(self.fiji_btn)
        main_left_box.Add(config_box,-1,wx.EXPAND)

        log_label=wx.StaticText(self,label="Log")
        log_label.SetFont(label_font)

        self.log_text = wx.TextCtrl(self,style = wx.TE_MULTILINE|wx.TE_RICH)
        #self.log_text.Disable()
        self.log_text.SetBackgroundColour(wx.LIGHT_GREY)
        bottom_left_box.Add(log_label,0,wx.EXPAND|wx.ALL)
        bottom_left_box.Add(self.log_text,-1,wx.EXPAND|wx.ALL)
        self.hub_label.SetFont(info_font)

        self.babel_hub = BabelBrushVolumeHub(None,self.write_message)
        self.index_button.Disable()
        self.del_button.Disable()
        self.launch_button.Disable()
        self.hub_label.SetLabel("None")

        self.volume_list = VolumeList(self,self.babel_hub)
        right_box.Add(self.volume_list,1,wx.EXPAND)
        self.SetSizer(outer_box)
        self.SetAutoLayout(True)
        self.Show(True)
        if base_folder:
            self.change_to_hub(base_folder)

    def launch_babel(self,event):
        exe_folder = os.path.realpath(os.path.dirname(sys.argv[0]))
        try:
            exe = os.path.join(exe_folder,"Babel Brush.exe")
            subprocess.Popen([exe,"--hub_folder",self.babel_hub.base_folder])
            sys.exit(0)
        except Exception as e:
            self.write_message("Unable to run {}".format(exe),True)

    def delete_volume(self,event):
        index  = self. volume_list.GetFirstSelected()
        self.babel_hub.delete_volume(index)
        self.volume_list.SetItemCount(len(self.volume_list.volume_hub.volumes))

    def addIndex(self,event):

        dlg = wx.FileDialog(
            self, message="Choose a file",
            defaultFile="",
            wildcard="*.nii",
            style=wx.FD_OPEN | wx.FD_CHANGE_DIR
            )
        if dlg.ShowModal() == wx.ID_OK:
            paths = dlg.GetPaths()
            try:
                self.babel_hub.index_nii_file(paths[0])
            except Exception as e:
                self.write_message("Cannot process {}. Is it .nii file".format(paths[0]))
                self.write_message(str(e),True)
                return

            self.volume_list.item_added()
            self.write_message("Indexed {}".format(paths[0]))

        dlg.Destroy()

    def create_new_hub(self,event):
        dlg = wx.DirDialog(
            self, message="Choose a Directory",
            style=wx.FD_OPEN
            )
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            try:
                create_directory_structure(path)
                self.write_message("Created directory structure")
            except Exception as e:
                self.write_message(str(e),True)
                self.write_message("Unable to create new hub",True)

            self.change_to_hub(path)

        dlg.Destroy()

    def write_message(self,msg,error=False):
        if error:
            self.log_text.SetDefaultStyle(wx.TextAttr(wx.RED,wx.LIGHT_GREY))
        self.log_text.AppendText(msg+"\n")
        self.log_text.SetDefaultStyle(wx.TextAttr(wx.BLACK,wx.LIGHT_GREY))

    def change_hub(self,event):
        dlg = wx.DirDialog(
            self, message="Choose a Directory",
            style=wx.FD_OPEN
            )
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            self.change_to_hub(path)

        dlg.Destroy()

    def change_fiji_path(self,event):
        dlg = wx.FileDialog(
            self, message="Choose fiji(imagej) exe",
            defaultFile="",
            wildcard="*.exe",
            style=wx.FD_OPEN | wx.FD_CHANGE_DIR
            )
        if dlg.ShowModal() == wx.ID_OK:
            paths = dlg.GetPaths()
            try:
                self.babel_hub.config["Fiji"]["path"]=paths[0]
                self.babel_hub.update_config()
                self.fiji_label.SetLabel(paths[0])
                self.write_message("Updated config with fiji path")
            except Exception as e:
                self.write_message("Unable to update config",True)
                self.write_message(str(e),True)

        dlg.Destroy()

    def change_to_hub(self,path):
        try:
            self.babel_hub.open_hub(path)
            self.volume_list.hub_changed()
            self.hub_label.SetLabel(path)
            self.index_button.Enable()
            self.del_button.Enable()
            self.launch_button.Enable()
            self.fiji_btn.Enable()
            fiji_path= self.babel_hub.config["Fiji"]["path"]
            if not fiji_path:
                fiji_path="none"
            self.fiji_label.SetLabel(fiji_path)
            self.write_message("Opened {}".format(path))

        except Exception as e:
            self.write_message("Could not open {}".format(path),True)
            self.write_message(str(e),True)

ex = wx.App()
BabelBrushVolumeManager(os.path.join(os.path.expanduser("~"),"Documents","Babel"))
ex.MainLoop()


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.pylab as pl
from matplotlib.colors import LogNorm, Normalize
import matplotlib as mpl
from astropy.io import fits, ascii
from astropy.modeling import models
from astropy import modeling
from astropy import units as u
# from kpicdrp import background, trace, extraction
import kpic_background as background
import kpic_trace as trace
import kpic_extraction as extraction
import os
from scipy import signal
import scipy.optimize as opt
import scipy
import multiprocessing as mp
from Nirspec_cmds import Spec_cmds
from Star_Tracker_cmds import Tracking_cmds
from Track_Cam_cmds import TC_cmds
from FAM_cmds import FAM_cmds
import Acquisition
import ktl
import Organization
import sys
import time
import warnings
import datetime
from time import sleep
from scipy.signal import medfilt as medfilt

'''
Required dependencies;
    kpic_background.py
    kpic_trace.py
    kpic_extraction.py
    hip81497_200602_sf2_wvsoln.txt
    2massK.txt
'''

### Get tons of these from extraction/rectification due to bad pixels
warnings.filterwarnings('ignore', category=RuntimeWarning)

# print('Current save directory:', Organization.get_path()[0])
cam   = TC_cmds()
track = Tracking_cmds()
FAM  = FAM_cmds()
spec  = Spec_cmds()
pd_sensitivity = 10
nreads = 10 #100

# connect to flip mirror keyword
serv = ktl.Service("kpic")
serv["PDPONAM"].monitor()

### These are some parameters that probably don't need to change often
N_order = 7
N_fiber = 1
nskip = 4 ### Only fit every 4 points for speed when using KPIC DRP, still give good enough trace after smoothing
itime = 1500 #30*1000
coadds = 1

def _get_cred2_flux(plot=False, nb_im=50):
    # return 1
    A = np.zeros(nb_im)
    for i in range(A.size):
        ### Get the processed cred2 image
        procim = track.Img.get_data(reform=True)
        ### Get location of psf
        goal = track.get_psf_goal()
        ### Check SHM tha tthese are right
        xcen = int(np.round(goal[1]))
        ycen = int(np.round(goal[2]))

        ### cutout - size tbd
        cutim = procim[ycen-50:ycen+51, xcen-50:xcen+51]
        if plot:
            plt.imshow(cutim, extent=[xcen-50, xcen+50, ycen-50, ycen+50])
            plt.show() ### plot to check

        totcnts = np.nansum(np.nansum(cutim))#/cutim.size
        ### Divide by itime to get flux
        A[i] = totcnts/cam.get_tint()
        # time.sleep(0.1)

    return np.mean(A)

def _acquire_fiber(fibr):
    '''
    Change to a new science fiber. If not given a valid fiber number, does nothing
    '''
    if fibr in [1,2,3,4]:
        track.set_goal("scf{}".format(fibr))
        # wait for tracking script to update goal
        while track.get_psf_goal()[3] != fibr:
            time.sleep(.1)

        start_t = time.time()
        # wait for a valid psf to be within .5 pixels in x and y of target
        while True:
            # calculate the dots to print based on how long we've been running
            dots = "."*((int(time.time()) - int(start_t))%3+1)
            # clear line
            sys.stdout.write(u"\u001b[2K")
            sys.stdout.write("\rfiber {} | acquiring".format(fibr)+dots)
            sys.stdout.flush()
            # get location of psf
            psfx, psfy = track.get_psf_cent()
            # get target location
            goal = track.get_psf_goal()

            # if we're close enough to goal, break
            if track.get_psf_parameters()[0] and abs(psfx - goal[1]) < .5 and abs(psfy - goal[2]) < .5:
                return 1
    else:
        return 0

def _rectify(cube, badpixmap, N_order, N_fiber, bkgd=np.zeros(0)):
    '''
    Flattens and rectifies a cube for peak finding
    '''
    ### Flatten and get rid of bad pixels
    deg = 4
    ### Badpixmap caussing issues, skipping for now
    # print('cube shape', cube.shape)
    # image = np.nansum(cube, axis=0) * badpixmap
    # Median instead of sum
    if bkgd.shape[0] == 0:
        image = np.nanmedian(cube, axis=0) * badpixmap
    else:
        image = (np.nanmedian(cube, axis=0) - bkgd) * badpixmap
    # plt.imshow(image, norm=LogNorm(vmin=1,vmax=1000))
    # plt.show()

    cind = int(image.shape[1]/2)
    nrows = image.shape[0]
    rowinds = np.arange(nrows)
    lind = cind-2
    rind = cind+2
    rectIm = np.empty(np.shape(image))
    rectIm[:,cind] = image[:,cind]
    # rectIm[:,cind] = np.nansum(image[:,cind-4:cind+5],axis=1) ### start from the middle, avoid edge effects
    rectCoeffs = np.empty((deg+1,image.shape[1]))
    init = np.zeros(deg+1)
    init[-2] = 1.
    rectCoeffs[:,cind] = init ### coeffs of the identity polynomial

    ### Find full peaks, then retain only the best N_orders*N_fibers
    if N_fiber==1:
        dist = 120
    else:
        dist = 15
    fullpeaks = signal.find_peaks(np.nansum(image[:,cind-4:cind+5],axis=1), prominence=1e2, distance=dist)
    peaks = fullpeaks[0]

    # plt.plot(rectIm[:,cind])
    # for peak in peaks:
    #   plt.axvline(peak)
    # plt.show()

    proms = fullpeaks[1]['prominences']
    while peaks.size > N_order*N_fiber:
        ### Toss partial traces at detector edge
        if min(peaks) < 40:
            worstpeak = np.argmin(peaks)
        else:
            worstpeak = np.argmin(proms)
        proms = np.delete(proms, worstpeak)
        peaks = np.delete(peaks, worstpeak)
    cpeaks = peaks
    ordcenters = [np.mean(cpeaks[i:i+N_fiber]) for i in range(0,N_order*N_fiber,N_fiber)]

    ### Loop over full image, in chunks
    rectfails = []
    fitinds = []
    while rind < image.shape[1] and lind > 0:
        ### Get left slice
        lcol = np.nanmedian(image[:,lind-4:lind+5], axis=1)
        lpeaks = signal.find_peaks(lcol, prominence=1e2, distance=dist)
        ### If missing peaks, skip
        if len(lpeaks[0]) < len(cpeaks):
            rectfails.append(lind)
            rectCoeffs[:,lind] = np.nan*np.ones(deg+1)
            rectIm[:,lind] = np.nan*np.ones(rowinds.size)
        else:
            peaks = lpeaks[0]
            ### If we get extra peaks, drop the weakest ones
            if len(peaks) > len(cpeaks):
                proms = lpeaks[1]['prominences']
                while len(peaks) > len(cpeaks):
                    ### Keep the N_fibers peaks closest to the center of each order, throw out rest
                    goodInds = np.empty(N_order*N_fiber)
                    for i in range(len(ordcenters)):
                        center = ordcenters[i]
                        inds = np.argsort([np.abs(peak-center) for peak in peaks])[:N_fiber]
                        goodInds[i*N_fiber:(i+1)*N_fiber] = inds
                    goodInds = np.sort(goodInds).astype(int)
                    peaks = peaks[goodInds]
                lpeaks = peaks
            else:
                lpeaks = lpeaks[0]
            ### Find polynomial to line up with reference peaks
            # print(lpeaks, cpeaks)
            lcoeffs = np.polyfit(lpeaks, cpeaks, deg=deg)
            rectCoeffs[:,lind] = lcoeffs
            lpoly = np.poly1d(lcoeffs)
            rectIm[:,lind] = np.interp(rowinds, lpoly(rowinds), lcol)

        ### Save polynomial coefficients
        fitinds.append(lind)
        lind-=4

        ### Get right slice
        rcol = np.nanmedian(image[:,rind-4:rind+5], axis=1)
        rpeaks = signal.find_peaks(rcol, prominence=1e2, distance=dist)
        ### If missing peaks, skip
        if len(rpeaks[0]) < len(cpeaks):
            rectfails.append(rind)
            rectCoeffs[:,rind] = np.nan*np.ones(deg+1)
            rectIm[:,rind] = np.nan*np.ones(rowinds.size)
        else:
            peaks = rpeaks[0]
            ### If extra peaks, drop weakest ones
            if len(peaks) > len(cpeaks):
                proms = rpeaks[1]['prominences']
                while len(peaks) > len(cpeaks):
                    ### Keep the N_fibers peaks closest to the center of each order, throw out rest
                    goodInds = np.empty(N_order*N_fiber)
                    for i in range(len(ordcenters)):
                        center = ordcenters[i]
                        inds = np.argsort([np.abs(peak-center) for peak in peaks])[:N_fiber]
                        goodInds[i*N_fiber:(i+1)*N_fiber] = inds
                    goodInds = np.sort(goodInds).astype(int)
                    peaks = peaks[goodInds]
                rpeaks = peaks
            else:
                rpeaks = rpeaks[0]
            ### Find polynomial to line uup with reference peaks
            rcoeffs = np.polyfit(rpeaks, cpeaks, deg=deg)
            rectCoeffs[:,rind] = rcoeffs
            rpoly = np.poly1d(rcoeffs)
            rectIm[:,rind] = np.interp(rowinds, rpoly(rowinds), rcol)
        ### Save polyomial coefficients
        fitinds.append(rind)
        rind+=4

    if len(rectfails) > 0:
        pass
        #print('Rectification failed in', len(rectfails), 'columns:', np.sort(rectfails))

    rectIm[np.isnan(rectIm)] = 0.0

    # plt.imshow(rectIm, norm=LogNorm(vmin=1,vmax=100))
    # plt.show()

    return rectIm, rectCoeffs, cpeaks

def _find_trace(rImage, rCoeffs, goodPeaks, N_order, N_fiber):
    '''
    Given a rectified image, rectification coefficients, peaks, and orders/fibers, finds the traces
    '''
    peaks = goodPeaks
    traces = np.empty((N_order*N_fiber, rCoeffs.shape[1]))
    for i in range(rCoeffs.shape[1]):
        ### Avoid weird things at the edges and bad pixels
        if np.any(np.isnan(rCoeffs[:,i])) or i < 20 or i > 2028:
            traces[:,i] = np.nan*np.ones(N_order*N_fiber)
            continue
        ### Undo the rectification
        for j in range(N_order*N_fiber):
            pcoeff = rCoeffs[:,i]
            pcoeff[np.abs(pcoeff)<1e-16] = 0.
            ### Solve the polynomial
            poly = np.poly1d(pcoeff)
            try:
                roots = (poly-peaks[j]).roots
            except:
                traces[j,i] = np.nan
                continue
            roots = np.real(roots[np.isreal(roots)])
            root = roots[roots>30] ### Reject bad answers near edges
            if root.size==1: traces[j,i] = float(root)
            elif root.size>1: traces[j,i] = np.min(root-peaks[j])
            else: traces[j,i] = np.nan
        ### Somehow these slip through
        if np.any(np.abs(traces[:,i]-peaks) > 100):
            traces[:,i] = np.nan*np.ones(N_order*N_fiber)
        ### Every now and fibers can switch
        for k in range(len(traces[:,i])-1):
            if traces[k,i] > traces[k+1,i]:
                traces[:,i] = np.nan*np.ones(N_order*N_fiber)
                continue
    ### Get rid of clear outliers that accumulate around edges
    traces[np.abs(traces)>2008] = np.nan
    traces[traces<40] = np.nan

    ### Polynomial nterpolate over NaNs
    for i in range(N_order*N_fiber):
        order = traces[i]
        inds = np.arange(order.size)
        ordfit = np.poly1d(np.polyfit(inds[~np.isnan(order)], order[~np.isnan(order)], deg=4))
        traces[i] = ordfit(inds)

    return traces

def _extract_fluxes(datacube, bkgd_noise, badpixmap, trace_locs, trace_sigs):
    pool = mp.Pool(4)
    flux_lst = []
    for i in range(len(datacube)):
        nframes = str(len(datacube))
        txt = '%3.0f / '%(i+1) + nframes
        sys.stdout.write('\rFlux extraction, frame'+txt)
        sys.stdout.flush()
        image = datacube[i]*badpixmap
        ### Get the diwth and location for each frame
        trace_sig = trace_sigs[i,:,:,:]
        trace_loc = trace_locs[i,:,:,:]

        ### This part was needed to get the flux extraction to work, it breaks if there's only one fiber
        ### illuminated (at least I think that's what's happening). Creates a dummy second fiber.
        if trace_loc.shape[0] == 1:
            arr = 20+trace_loc
            trace_loc = np.append(trace_loc, arr,axis=0)
            trace_sig = np.append(trace_sig, trace_sig, axis=0)
        ### Background traces for DRP
        trace_loc_slit, trace_loc_dark = trace.get_background_traces(trace_loc)
        ### More DRP setup
        trace_loc_wbkg = np.concatenate([trace_loc, trace_loc_slit, trace_loc_dark], axis=0)
        trace_sigmas_wbkg = np.concatenate([trace_sig, trace_sig, trace_sig], axis=0)
        trace_flags = np.array([0, ] * trace_loc.shape[0] + [1, ] * trace_loc.shape[0] + [2, ] * trace_loc.shape[0])
        ### DRP frame processing
        fluxes, errors_extraction, errors_bkgd_only = extraction.extract_flux(image, trace_loc_wbkg, trace_sigmas_wbkg,
                                                                          output_filename=None,
                                                                          img_hdr=None,
                                                                          img_noise=bkgd_noise, fit_background=True,
                                                                          trace_flags=trace_flags, pool=pool,
                                                                          bad_pixel_fraction=0.01,
                                                                          box=True)
        flux_lst.append(fluxes)
    print('')
    int_fluxes = np.nansum(flux_lst, axis=-1)
    ### If averaging < 100 counts/pixel, more likely to run into trace fitting issues
    #if np.nanmax(int_fluxes) < 100*N_order*2048:
        #print('Warining: Low counts on spec, trace fitting may be unreliable. Check source intensity!')
    return np.asarray(flux_lst)

def _extract_fluxes_hybrid(datacube, bkgd_noise, badpixmap, trace_loc):
    '''
    Does DRP flux extraction using trace locations from the faster trace finding code
    '''
    trace_sigs = np.empty((len(datacube), N_fiber, N_order, 2048))
    if datacube.ndim == 3:
        image = np.nansum(datacube, axis=0)
    else: image = datacube

    image[np.isnan(image)] = 0.0
    image[np.isnan(image)] = 0.0
    _, _, trace_sig, _ =_fit_traces_fast(image, trace_loc, N_order, N_fiber)

    flux_lst = []
    pool = mp.Pool(4)
    for i, image in enumerate(datacube):
        nframes = str(len(datacube))
        txt = '%3.0f / '%(i+1) + nframes
        sys.stdout.write('\rExtracting frame'+txt)
        sys.stdout.flush()

        if trace_loc.shape[0] == 1:
            arr = 20+trace_loc
            trace_loc = np.append(trace_loc, arr,axis=0)
            trace_sig = np.append(trace_sig, trace_sig, axis=0)
        ### Background traces for DRP
        trace_loc_slit, trace_loc_dark = trace.get_background_traces(trace_loc)
        ### More DRP setup
        trace_loc_wbkg = np.concatenate([trace_loc, trace_loc_slit, trace_loc_dark], axis=0)
        trace_sigmas_wbkg = np.concatenate([trace_sig, trace_sig, trace_sig], axis=0)
        trace_flags = np.array([0, ] * trace_loc.shape[0] + [1, ] * trace_loc.shape[0] + [2, ] * trace_loc.shape[0])
        ### DRP frame processing
        fluxes, errors_extraction, errors_bkgd_only = extraction.extract_flux(image, trace_loc_wbkg, trace_sigmas_wbkg,
                                                                          output_filename=None,
                                                                          img_hdr=None,
                                                                          img_noise=bkgd_noise, fit_background=True,
                                                                          trace_flags=trace_flags, pool=pool,
                                                                          bad_pixel_fraction=0.01,
                                                                          box=True)
        flux_lst.append(fluxes)
    print('')
    return np.asarray(flux_lst)

def _extract_fluxes_fast(datacube, trace_loc, N_order=7, bpmap=np.ones((2048,2048))):
    trace_sigs = np.empty((len(datacube), N_fiber, N_order, 2048))
    flux_ests = np.empty((len(datacube), N_fiber, N_order, 2048))
    tot_cnts = np.empty((len(datacube), N_fiber, N_order, 2048))
    trace_locs = np.empty((len(datacube), N_fiber, N_order, 2048))

    # TO-DO: look into parallizing this to speed up
    for i, image in enumerate(datacube):
        # nframes = str(len(datacube))
        # txt = '%3.0f / '%(i+1) + nframes
        # sys.stdout.write('\rExtracting frame'+txt)
        # sys.stdout.flush()

        # apply bad pixel map
        goodim = image * bpmap
        goodim[np.isnan(goodim)] = 0.0
        trace_amp, _, trace_sig, tot_cnt =_fit_traces_fast(goodim, trace_loc, N_order, N_fiber)
        # trace_locs[i] = trace_loc
        # flux_ests[i] = trace_sig*trace_amp*np.sqrt(2*np.pi)
        # flux_ests[flux_ests>3e5] = 0.0 ### Reject hot pixel values
        # flux_ests[flux_ests<0] = 0.0 ### Reject bad pixel values\

        trace_sigs[i] = trace_sig
        tot_cnts[i] = tot_cnt

    # print('')
    return tot_cnts

def _fit_traces(datacube, badpixcube, fibers, nskip=4):
    '''
    Use the KPIC DRP to fit each trace for each image in cube. Requires
    a good guess for fiber positions. Nskip determines how many pixels to skip
    in the fitting, setting it higher speeds things up.
    '''
    # print('badpixshape', badpixcube.shape)
    # plt.imshow(badpixcube[0])
    # plt.show()
    nthreads = 4
    ### Get number of orders from fibers
    norder = len(fibers[list(fibers.keys())[0]])
    _, nx, ny = badpixcube.shape
    ### Fit each trace in the cube
    trace_list = []
    for i in range(len(datacube)):
        nframes = str(len(datacube))
        txt = '%3.0f / '%(i+1) + nframes
        sys.stdout.write('\rTrace fitting, frame'+txt)
        sys.stdout.flush()
        ### Setup for DRP tools - expecting list of images
        image = np.expand_dims(datacube[i],axis=0)
        fiber_list = trace.guess_star_fiber(image, fibers)

        ### Fit traces and smooth using KPIC DRP tools
        trace_calib, residuals = trace.fit_trace(fibers,fiber_list,image[:,:,::nskip],badpixcube[:,:,::nskip],ny,int(nx/nskip),norder,numthreads=nthreads, fitbackground=False)
        # polyfit_trace_calib, smooth_trace_calib = trace.smooth(trace_calib)

        ### Save outputs to arrays for plotting
        trace_list.append(trace_calib)
    print('')
    ### Get widths and locations
    trace_list = np.asarray(trace_list)
    # print(trace_list.shape)
    sigs = trace_list[:,:,:,:,1]
    # print(sigs.shape)
    # print(sigs[0,0,1])
    for i in range(len(sigs)):
        try:
            sigs[i] = _fourier_smooth(sigs[i], trace_list.shape[2], trace_list.shape[1])
        except:
            sigs[i] = np.nan*np.ones(np.shape(sigs[i]))
            print('Bad frame', i)

    ### Resample onto full grid if needed
    if sigs.shape[3] != 2048:
        oldshape = sigs.shape
        oldlen = sigs.shape[3]
        oldgrid = np.arange(oldlen)
        newshape = (oldshape[0], oldshape[1], oldshape[2], 2048)
        new_trace_sig = np.empty(newshape)
        newgrid = np.linspace(0, oldlen, num=2048)
        for i in range(newshape[0]):
            for j in range(newshape[1]):
                for k in range(newshape[2]):
                    new_trace_sig[i,j,k] = np.interp(newgrid, oldgrid, sigs[i,j,k])
        sigs = new_trace_sig

    return sigs

def _one_order_fit(traces, goodim, i, j):
    imslice = np.empty((9,2048))
    amps = np.empty(2048)
    for k in range(2048):
        tcen = int(np.round(traces[i,j,k],0))
        cut = np.nansum(goodim[tcen-4:tcen+5,k:k+1], axis=1)

        ### Avoid wackiness at detector edge
        if cut.size != 9 or np.abs(k-1024) > 1020:
            imslice[:,k] = np.nan*np.ones(9)
            amps[k] = np.nan
        else:
            # bkg = np.nanmedian(np.concatenate((cut[:5], cut[:-5])))
            imslice[:,k] = cut #-bkg
            ### Deal with sub-pixel phase by fitting peak
            p1 = np.polyfit(np.arange(3,6), cut[3:6], 2)
            pfit = np.poly1d(p1)
            fitprof = pfit(np.arange(3,6.1,0.1))
            amps[k] = np.nanmax(fitprof)
            # plt.plot(cut)
            # plt.plot(np.arange(3,6.1,0.1), fitprof)
            # plt.show()

    ### Because it's a sampled gaussian, amps will be underestimate on average
    ### Leads to sigs being slightly larger - underestimate resolution
    pshifts = np.abs(np.argmax(imslice, axis=0)-4)
    ### Need to figure out Lorentzian broadening here, thinking it's ~5 percent
    tots = np.nansum(imslice, axis=0)
    sigs = tots/(amps*np.sqrt(2*np.pi))
    amps[pshifts>1.5] = np.nan
    sigs[pshifts>1.5] = np.nan
    sigs[sigs>3] = np.nan

    return amps, sigs, tots

# def _fit_traces_fast_new(goodim, traces, N_order, N_fiber):
#     rawsig = np.empty((N_fiber, N_order, 2048))
#     rawamp = np.empty((N_fiber, N_order, 2048))
#     rawtot = np.empty((N_fiber, N_order, 2048))

#     if traces.ndim == 2:
#         traces = np.expand_dims(traces,axis=0)

#     for i in range(N_fiber):
#         # parallelize N_order calc
#         pool = mp.Pool(processes=N_order)
#         res_list = []
#         for j in range(N_order):
#             res_list.append(pool.apply_async(_one_order_fit, args=(traces, goodim, i, j)))

#         for j in range(N_order):
#             res = res_list[j].get()
#             rawamp[i,j] = res[0]
#             rawsig[i,j] = res[1]
#             rawtot[i,j] = res[2]

#         pool.close()
#     # print(rawsig)
#     ### Fourier smoothing - drop all the high-frequency stuff
#     # smoothamp = _fourier_smooth(rawamp, N_order, N_fiber)
#     smoothsig = _fourier_smooth(rawsig, N_order, N_fiber)
#     return rawamp, rawsig, smoothsig, rawtot

def _fit_traces_fast(goodim, traces, N_order, N_fiber):
    rawsig = np.empty((N_fiber, N_order, 2048))
    rawamp = np.empty((N_fiber, N_order, 2048))
    rawtot = np.empty((N_fiber, N_order, 2048))

    if traces.ndim == 2:
        traces = np.expand_dims(traces,axis=0)

    for i in range(N_fiber):
        for j in range(N_order):
            imslice = np.empty((9,2048))
            amps = np.empty(2048)
            for k in range(2048):
                tcen = int(np.round(traces[i,j,k],0))
                cut = np.nansum(goodim[tcen-4:tcen+5,k:k+1], axis=1)

                ### Avoid wackiness at detector edge
                if cut.size != 9 or np.abs(k-1024) > 1020:
                    imslice[:,k] = np.nan*np.ones(9)
                    amps[k] = np.nan
                else:
                    # bkg = np.nanmedian(np.concatenate((cut[:5], cut[:-5])))
                    imslice[:,k] = cut #-bkg
                    ### Deal with sub-pixel phase by fitting peak
                    p1 = np.polyfit(np.arange(3,6), cut[3:6], 2)
                    pfit = np.poly1d(p1)
                    fitprof = pfit(np.arange(3,6.1,0.1))
                    amps[k] = np.nanmax(fitprof)
                    # plt.plot(cut)
                    # plt.plot(np.arange(3,6.1,0.1), fitprof)
                    # plt.show()

            ### Because it's a sampled gaussian, amps will be underestimate on average
            ### Leads to sigs being slightly larger - underestimate resolution
            pshifts = np.abs(np.argmax(imslice, axis=0)-4)
            ### Need to figure out Lorentzian broadening here, thinking it's ~5 percent
            tots = np.nansum(imslice, axis=0)
            sigs = tots/(amps*np.sqrt(2*np.pi))
            amps[pshifts>1.5] = np.nan
            sigs[pshifts>1.5] = np.nan
            sigs[sigs>3] = np.nan
            rawamp[i,j] = amps
            rawsig[i,j] = sigs
            rawtot[i,j] = tots
    # print(rawsig)
    ### Fourier smoothing - drop all the high-frequency stuff
    # smoothamp = _fourier_smooth(rawamp, N_order, N_fiber)
    smoothsig = _fourier_smooth(rawsig, N_order, N_fiber)

    return rawamp, rawsig, smoothsig, rawtot

def _fourier_smooth(arr, N_order, N_fiber, cutoff=0.004):
    '''
    Fourier-based smoothing for widths
    '''
    arr[arr<0.5] = np.nan
    smootharr = np.empty(arr.shape)
    try:
        for i in range(N_fiber):

            for j in range(N_order):
                # res = res_list[j].get()
                # inds = res[0]
                # goodinds = res[1]
                # newvals = res[2]
                inds = np.arange(arr[i,j].size)
                vals =  arr[i,j]
                badinds = np.isnan(vals)
                goodinds = inds[~badinds]
                goodvals = vals[~badinds]
                # pdb.set_trace()
                ft = np.fft.fft(goodvals)
                freq = np.fft.fftfreq(goodinds.shape[-1])
                ### This has been tuned manually based on fringe patterns
                ### TODO - replace with a proper window to reduce ringing
                ft.imag[np.abs(freq)>cutoff] = 0.
                ft.real[np.abs(freq)>cutoff] = 0.
                newvals = np.fft.ifft(ft)
                smootharr[i,j] = np.interp(inds, goodinds, newvals)
                # plt.plot(vals)
                # plt.plot(smootharr[i,j])
                # plt.show()
        return smootharr
    except ValueError:
        print('Smoothing failed')
        return arr

def set_nspec_reads(tint):

    if 1400 <= tint < 3000:
        n_reads = 1
    elif 3000 <= tint < 4500:
        n_reads = 2
    elif 4500 <= tint < 6000:
        n_reads = 3
    elif 6000 <= tint <= 60000:
        n_reads = 4
    else:
        n_reads = 16
    # CDS is 2
    if n_reads == 1 and not ktl.read('nspec','sampmode') == 'CDS':
        ktl.write('nspec','sampmode', '2')
    # MCDS is 2
    elif n_reads > 1 and not ktl.read('nspec','sampmode') == 'MCDS':
        ktl.write('nspec','sampmode', '3')

    ktl.write('nspec','numreads', str(n_reads))
    time.sleep(0.1)

def _take_backgrounds(itime, coadds, nframes=5, retake=False, Path=None):
    '''
    Take background frames for scan
    '''
    fam_bkgd_pos = [50, 9900]

    ### Path info
    if Path is None:
        Path = Organization.get_path('Spec')
        Path = list(Path)
        Path[1]=Path[1].split('.')[0]

    # Path is actually a list, so get the string
    real_path = Path[0]
    #print(real_path)

    ### Check that backgrounds already exist, if they don't force a retake
    bkg_names = []
    for i in range(nframes):
        bkg_names.append(real_path + 'bkgd_'+str(i)+'_itime'+str(itime)+'.fits')
        if not os.path.isfile(bkg_names[-1]):
            retake = True
    track_flag=False
    if retake:
        # set nreads and sampmode
        set_nspec_reads(itime)

        ### Stop tracking and move off-slit for backgrounds
        if track.is_tracking():
            track_flag=True
            track.stop_tracking()
            # Wait for the Tracking loop to complete its last iteration
            time.sleep(2)
        ### Get initial TTM position
        TTM_CP = FAM.get_pos()
        ### Move TTM and wait for new position
        FAM.set_pos(fam_bkgd_pos, block=True)

        ### Take backgrounds
        for i in range(nframes):
            # Take backgrounds, save for DRP
            spec.TakeImage(Coadds=coadds, tint=itime, Save=True, filename='bkgd_'+str(i)+'_itime'+str(itime), checkbadframe=False)
        FAM.set_pos([TTM_CP[0], TTM_CP[1]], block=True)

    ### Return to original TTM position and restart tracking
    bkgd, smoothed_thermal_noise, badpixmap = background.make_badpixmap(bkg_names)

    if track_flag==True:
        track.start_tracking()

    return bkgd, smoothed_thermal_noise, badpixmap

# Make trace calc separate from flux extract
def _calc_trace(datacube, badpixmap, N_order, N_fiber, bkgd=np.zeros(0), plot=False):

    # datacube should have be numpy array of shape (Nframes, x, y)
    # Rectify images
    rectIm, rectCoeffs, goodPeaks = _rectify(datacube, badpixmap, N_order, N_fiber, bkgd=bkgd)
    trace_locs = _find_trace(rectIm, rectCoeffs, goodPeaks, N_order, N_fiber)

    if plot:
        collapsedim = np.nansum(datacube,axis=0)
        collapsedim[np.isnan(collapsedim)] = 0.
        plt.title('Stacked cube, with traces')
        plt.imshow(collapsedim, norm=LogNorm(vmin=1,vmax=100))
        for ordloc in trace_locs:
          plt.plot(ordloc, color='r', linestyle='--')
        plt.show()

    return trace_locs

def _total_flux_frame(this_image, badpixmap, N_order, N_fiber, bkgd=np.zeros(0), loaded_Trace=np.zeros(0), which_Order=0):
    '''
    Get the total flux on spec for a single image
    '''

    datacube = np.zeros((1, this_image.shape[0], this_image.shape[1]))

    if bkgd.shape[0] == 0:
        datacube[0] = this_image
    else:
        datacube[0] = this_image - bkgd    # remove a master bkgd

    # print('here0')

    # Trace fitting if not given one
    if loaded_Trace.shape[0] == 0:
        # Rectify images
        rectIm, rectCoeffs, goodPeaks = _rectify(datacube, badpixmap, N_order, N_fiber)
        trace_locs = _find_trace(rectIm, rectCoeffs, goodPeaks, N_order, N_fiber)
    else:
        trace_locs = loaded_Trace
        # print('Using input trace locations.')

    # Fast flux extractiong
    if N_fiber == 1:
        trace_locs = np.expand_dims(trace_locs, axis=0)
    # print('here1')

    # flux_lst has shape (Nframes, Nfibers, Norder, 2048)
    flux_lst =_extract_fluxes_fast(datacube, trace_locs, N_order=N_order, bpmap=badpixmap)

    order_flux = np.nansum(flux_lst[0][0][which_Order])

    # print('here2')

    # Get total flux in each frame
    int_flux = np.empty(flux_lst.shape[0])
    for i, frame in enumerate(flux_lst):
        int_flux[i] = np.nansum(np.nansum(frame, axis=2), axis=1)[0]

    return int_flux, order_flux

def _total_flux_framepath(frame_path, badpixmap, N_order, N_fiber, bkgd=np.zeros(0), loaded_Trace=np.zeros(0), which_Order=0):
    '''
    Get the total flux on spec for each image in the cube. Moved here since it's used a bunch
    '''

    ### Load the frame into a datacube
    hdu = fits.open(frame_path)
    this_image = hdu[0].data.T[:,::-1]
    datacube = np.zeros((1, this_image.shape[0], this_image.shape[1]))

    if bkgd.shape[0] == 0:
        datacube[0] = this_image
    else:
        datacube[0] = this_image - bkgd    # remove a master bkgd

    ### Trace fitting if not given one
    if loaded_Trace.shape[0] == 0:
        rectIm, rectCoeffs, goodPeaks = _rectify(datacube, badpixmap, N_order, N_fiber)
        trace_locs = _find_trace(rectIm, rectCoeffs, goodPeaks, N_order, N_fiber)
    else:
        trace_locs = loaded_Trace
        # print('Using input trace locations.')

    # print(trace_locs.shape)
    ### Fast flux extractiong
    if N_fiber == 1:
        trace_locs = np.expand_dims(trace_locs, axis=0)

    # flux_lst has shape (Nframes, Nfibers, Norder, 2048)
    flux_lst =_extract_fluxes_fast(datacube, trace_locs, N_order=N_order, bpmap=badpixmap)

    order_flux = np.nansum(flux_lst[0][0][which_Order])

    ### Get total flux in each frame
    int_flux = np.empty(flux_lst.shape[0])
    for i, frame in enumerate(flux_lst):
        int_flux[i] = np.nansum(np.nansum(frame, axis=2), axis=1)[0]
    ### If averaging < 100 counts/pixel, more likely to run into trace fitting issues
    #if np.nanmax(int_flux) < 100*N_order*2048:
        #print('Warining: Low counts on spec, trace fitting may be unreliable. Check source intensity!')

    return int_flux, order_flux, trace_locs

def _reduce_drp_byframe(datacube, badpixmap, smoothed_thermal_noise,  N_order, N_fiber, stack=False):
    '''
    Get trace widths and fluxes for each frame in the datacube, using KPIC DRP methods (slow)
    '''
    ### Really wants a cube for this

    # badpixmap = np.ones((2048,2048))
    badpixmap = np.ones((2048,2048))
    badpixcube = np.expand_dims(badpixmap, axis=0)
    ### Rectify frame and find trace locations
    # print('Datacube shape', datacube.shape)
    rectIm, rectCoeffs, goodPeaks = _rectify(datacube, badpixmap, N_order, N_fiber)
    ### Check rectified image if desired
    # plt.imshow(rectIm, norm=LogNorm(vmin=1, vmax=100))
    # plt.show()
    traces = _find_trace(rectIm, rectCoeffs, goodPeaks, N_order, N_fiber)
    trace_locs = np.empty((len(datacube), N_fiber, N_order, 2048))
    for i in range(len(datacube)):
        trace_locs[i] = traces

    ### Trace fitting. Currently using KPIC DRP (slow)
    ### This is for the current KPIC DRP trace fitting process
    tracebounds = []
    for order in traces:
        tracebounds.append([int(np.min(order)-20), int(np.max(order)+20)])
    ### Assign to fibers for the DRP
    fibers = {}
    for i in range(N_fiber):
        fibers[i] = tracebounds[i::N_fiber]
    ### Fit trace widths with KPIC DRP
    if not stack:
        trace_sigs = _fit_traces(datacube, badpixcube, fibers, nskip=nskip)
    else:
        trace_sigs = []
        trace_sig = _fit_traces(np.expand_dims(np.nansum(datacube,axis=0),axis=0), badpixcube, fibers, nskip=nskip)
        for i in range(len(datacube)):
            trace_sigs.append(trace_sig[0])
        trace_sigs = np.asarray(trace_sigs)
    # ## Do flux extraction with KPIC DRP tools, box option
    flux_lst = _extract_fluxes(datacube, smoothed_thermal_noise, badpixmap, trace_locs, trace_sigs)

    return trace_sigs, trace_locs, flux_lst


db_path = r".\db\test.db"

from flask import Flask, render_template, request
# from flask_bootstrap import Bootstrap
from Database import Database
import os
import hashlib
from PIL import Image
import json

app = Flask(__name__, template_folder =  './assets')

def readFromSBP(plate_id):
    try:
        hash_code = "0"
        status = "NOT OK"
        with open("example.json", encoding='utf-8') as f:
            data = json.load(f)
            hash_code = data[plate_id]["hash"]
            status = data[plate_id]["status"]
        return hash_code, status
    except Exception as e:
        print(e)
        return "0", "NOT OK"

@app.route('/', methods = ['GET', 'POST'])
def homePage():
    hash_key = "0"
    filename = "images/template.png"
    if request.method == 'POST':
        try:
            db = Database(db_path)
            hash_key, photo_path = db.selectLast()
            img = Image.open(photo_path)
            img.save("./static/images/car_image.png", "PNG")
            filename = 'images/car_image.png'
            # photo_path = photo_path.replace("\\", "/")
        except:
            print("Couldn't connect to database or load image")

        first_two = hash_key[:2]
        last_two = hash_key[-2:]
        hash_key = hash_key

    return render_template("index.html", hash_key = hash_key, filename = filename)

@app.route('/input', methods = ['GET' , 'POST'])
def inputPage():
    hide = 1
    new_hash_key = "0"
    hash_key = "0"
    filename = "images/template.png"
    if(request.method == 'POST'):
        iin = request.form['iin']
        plate_id = request.form['plate_id']
        print(plate_id)
        new_hash_key = hashlib.sha256(plate_id.encode()).hexdigest()
        print(new_hash_key)
        db = Database(db_path)
        hash_key, photo_path = db.selectByHash(new_hash_key)
        if(hash_key != None):
            img = Image.open(photo_path)
            img.save("./static/images/input_car_image.png", "PNG")
            filename = 'images/input_car_image.png'
            hide = 0
        else:
            hide = 2
            hash_key = "0"

    return render_template("input.html", hide = hide, new_hash_key = new_hash_key, hash_key = hash_key, filename = filename)

@app.route('/input-v2', methods = ['GET' , 'POST'])
def inputV2Page():
    hide = 1
    new_hash_key = "0"
    hash_key = "0"
    filename = "images/template.png"
    status = "NOT OK"
    if(request.method == 'POST'):
        iin = request.form['iin']
        plate_id = request.form['plate_id']
        print(plate_id)
        new_hash_key = hashlib.sha256(plate_id.encode()).hexdigest()
        print(new_hash_key)
        db = Database(db_path)
        new_hash_key, photo_path = db.selectByHash(new_hash_key)
        hash_key, status = readFromSBP(plate_id)

        if(new_hash_key != None):
            img = Image.open(photo_path)
            img.save("./static/images/input_car_image.png", "PNG")
            filename = 'images/input_car_image.png'
            hide = 0
        else:
            new_hash_key = ""
            hide = 2

    return render_template("input_v2.html", hide = hide, status = status ,new_hash_key = new_hash_key, hash_key = hash_key, filename = filename)

if __name__ == '__main__':
    app.run(host = "localhost", port = 5000)


"""
1-100 arasında rastgele üretilecek bir sayıyı aşağı yukarı ifadeleri ile buldurmaya çalışın
** "random modülü" için "python random" şekilde arama yapın
** 100 üzerinde puanlama yapın. Her soru 20 puan.
** Hak bilgisini kullanıcıdan alın ve her soru belirtilen can sayısı üzerinden hesaplansın.

- Oyun hakkında :
Tahmin edilen sayıyı soracak. Tahmin büyükse : aşağı, küçükse : yukarı gibi yönlendirmeler yapacak
Hak = 5(varsayılan olarak) Her bilemediği hakda puanı 100 üzerinden 20 puan aşağı inecek
Daha sonraki aşamada Hak bilgisi kullanıcı tarafından belirlenecek. 10 hak ise her soru 10 puan
"""

import random

rnd = random.randint(1,100)
can = int(input("Kaç hakkınız olsun : "))
hak = can
sayac = 0
puan = 100
while hak > 0:
    hak -= 1
    sayac += 1
    tahmin = int(input("Tahmin : "))
    if (rnd == tahmin):
        print(f"Tebrikler sayıyı {sayac}. defada bildiniz. Puanınız : {100-(100/can)*(sayac-1)}")
        break
    elif (rnd > tahmin):
        print("Daha büyük bir sayı deneyin.")
    else:
        print("Daha küçük bir sayı deneyin.")

    if hak == 0:
        print(f"Tahmin hakkınız bitti! Tutulan sayı = {rnd}")


from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Ridge
from sklearn.tree import DecisionTreeRegressor
from sklearn import linear_model
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestRegressor
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.naive_bayes import GaussianNB
from sklearn import svm
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn import tree
import pickle
import json
import re
from datetime import datetime
from pathlib import Path
from sklearn.metrics import accuracy_score
from collections import OrderedDict
dummy=pd.DataFrame()
X_train=pd.DataFrame()
y_train=pd.DataFrame()
X_test=pd.DataFrame()
y_test=pd.DataFrame()
label_encoder = LabelEncoder()
label_encoder1=LabelEncoder()
label_encoderx=LabelEncoder()
filepath1=[]
feature_count=""
sel_x=()
test_df=pd.DataFrame()
target_df=pd.DataFrame()
comman_list=[]
tt=[]
RANDOM_STATE=10
def val(n):
    if n == 0:
        print("entered into if")
        json_fail={"status":"false","msg":"0 rows after preprocessing ,choose target and x features carefully"}
        result_json=json.dumps(json_fail)
        print(result_json)
        return result_json
    else:
        return 'continue'
def hint(df,x,y):
    record=list(df.iloc[0])
    ex_dict=OrderedDict()
    for i in range(len(record)):
        ex_dict[x[i]]=str(record[i])
    print(ex_dict)
    ex_dict.pop(y, None)
    print(ex_dict)
    return ex_dict
def preprocess_dropna2(filepath,ex,x,y,algo,tt_split,mid,algo_params,email,label_encoder_file_path):
    global dummy,filepath1,feature_count,sel_x,target_df,comman_list,tt
    tt=tt_split
    print('came in',label_encoder_file_path)
    sel_x=x
    #print(sel_x)
    feature_count=int(len(x))
    #print(feature_count)
    filepath1=filepath
    if ex in ['.csv']:
        try:
            df_preprocess=pd.read_csv(filepath,encoding='raw_unicode_escape')
        except UnicodeDecodeError:
            failure={"status":"false","msg":"uploaded file cannot be used for modelling pls upload proper file"}
            success_json=json.dumps(failure)
            return success_json

        list_x=x
       # print(list_x)
        list_y=str(y)
        target_df=pd.DataFrame({'target':df_preprocess[list_y]})
        comman_list=x
        comman_list.append(list_y)
        print(comman_list)
        required_df=df_preprocess[comman_list]
        required_df=required_df.dropna(axis=0)
        shape=required_df.shape
        print(shape)
        zero=shape[0]
        print(zero)
        res=val(zero)
        if res!='continue':
            return res

       # print(res)
        #return res
        else:
            ex_dict=hint(required_df,x,y)
            dummy=dummy.append(required_df)
            visual()
            print('yet to encode')
            encoding(required_df,filepath,label_encoder_file_path)
        #normalize()
        #print(list_x)
        #print(list_y)
            print('encoding done')
            x1=dummy.iloc[:,0:feature_count]
            y1=dummy.iloc[:,-1]
            print(x1)
            data_spilt(x1,y1)
            result_json=model(algo,ex_dict,mid,algo_params,email)
            return result_json
    else:
        df_preprocess=pd.read_excel(filepath,encoding='latin')
        list_x=x
        list_y=y
        comman_list=list_x
        comman_list.append(list_y)
        print(comman_list)
        required_df=df_preprocess[comman_list]
        required_df=required_df.dropna()
        print(required_df.shape)
        shape=required_df.shape
        zero=shape[0]
        res=val(zero)
        if res!='continue':
            return res
        else:
            ex_dict=hint(required_df,x,y)
            dummy=dummy.append(required_df)
            visual()
            encoding(required_df,filepath,label_encoder_file_path)
            #normalize()
            print(target_df)
            x1=dummy.iloc[:,0:feature_count]
            y1=dummy.iloc[:,-1]
            data_spilt(x1,y1)
            result_json=model(algo,ex_dict,mid,algo_params,email)
            return result_json
def visual():
    print(dummy)
def encoder_using_path(le,df_x):
    for i in df_x.select_dtypes(['object']):
        if i in le:
            temp =le.get(i)
            # my_dict2= {y:x for x,y in temp.items()}
            # label[i]=dict(my_dict2)
            df_x[i] = df_x[i].map(temp)
            print("encoded" ,i)
    return df_x
def encoding(required_df,filepath,label_encoder_file_path):
    global dummy
    print(label_encoder_file_path)
    with open(label_encoder_file_path, 'rb') as file:
        le=pickle.load(file)
    req_df=encoder_using_path(le,required_df)
    dummy=req_df
def normalize():
    global dummy
    scaler = StandardScaler()
    scaled_features=scaler.fit_transform(dummy)
    scaled_features_df = pd.DataFrame(scaled_features, index=dummy.index, columns=dummy.columns)
    dummy =scaled_features_df
    visual()
def data_spilt(x,y):
    print("test split",x)
    global X_train,X_test,y_train,y_test,tt
    if "80% train and 20% test" in tt:
        X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.20)
        print(X_train.shape,y_train.shape,X_test.shape,y_test.shape)
    elif "70% train and 30% test" in tt:
        X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.30)
        print(X_train.shape,y_train.shape,X_test.shape,y_test.shape)
    elif "50% train and 50% test" in tt:
        X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.50)
        print(X_train.shape,y_train.shape,X_test.shape,y_test.shape)
def algo_select(li):
    global X_train,y_train,X_test,y_test
    # for i in li:
    #     if "DecisionTreeClassifier" in i:
    #         para=[]
    #         para_value=[]
    #         for j in li:
    #             par=li[j]
    #             for p in par:
    #                 para_value.append(par.get(p))
    #             clf = tree.DecisionTreeClassifier(p=v)
    #             clf = clf.fit(X_train,y_train)
    #             pre=clf.predict(X_test)
    #             acc=accuracy_score(y_test,pre)
    #             print("training done")
    #             return acc*100
    for u in li:
        algo_tune=u
        if "DecisionTreeClassifier" in algo_tune:
            var=[]
            keys=[]
            for key,value in li[u].items():
                if value !="":
                    var.append(key)
                    keys.append(value)
            print(var,keys)
            if "max_depth" in var:
                indi=var.index('max_depth')
                max_depths=int(keys[indi])
            else:
                max_depths=None
            if "max_features" in var:
                indi1=var.index('max_features')
                max_featuress=keys[indi1]
            else:
                max_featuress=None
            if "min_samples_leaf" in var:
                indi2=var.index('min_samples_leaf')
                min_samples_leafs=int(keys[indi2])
            else:
                min_samples_leafs=1
            if "criterion" in var:
                indi3=var.index('criterion')
                criterions=keys[indi3]
            else:
                criterions='gini'
            if "min_samples_split" in var:
                indi4=var.index('min_samples_split')
                min_samples_splits=int(keys[indi4])
            else:
                min_samples_splits=2
            clf = tree.DecisionTreeClassifier(max_depth=max_depths,max_features=max_featuress,min_samples_leaf=min_samples_leafs,criterion=criterions,min_samples_split=min_samples_splits,random_state=RANDOM_STATE)
            clf = clf.fit(X_train,y_train)
            pre=clf.predict(X_test)
            acc=accuracy_score(y_test,pre)
            print("training done",acc)
            return acc
        if "RandomForestClassifier" in algo_tune:
            var1=[]
            keys1=[]
            for key,value in li[u].items():
                if value !="":
                    var1.append(key)
                    keys1.append(value)
            if "max_depth" in var1:
                indi8=var1.index('max_depth')
                max_depths=keys1[indi8]
            else:
                max_depths=None
            if "n_estimators" in var1:
                indi5=var1.index('n_estimators')
                n_estimatorss=int(keys1[indi5])
            else:
                n_estimatorss=10
            if "min_samples_leaf" in var1:
                indi6=var1.index('min_samples_leaf')
                min_samples_leafs=int(keys1[indi6])
            else:
                min_samples_leafs=1
            if "min_samples_split" in var1:
                indi7=var1.index('min_samples_split')
                min_samples_splits=int(keys1[indi7])
            else:
                min_samples_splits=2
            if "criterion" in var1:
                indi9=var1.index('criterion')
                criterions=keys1[indi9]
            else:
                criterions='gini'
            if "max_features" in var1:
                indi10=var1.index('max_features')
                max_featuress=keys1[indi10]
            else:
                max_featuress=None
            clf1=RandomForestClassifier(n_estimators=n_estimatorss,max_depth=max_depths,min_samples_split=min_samples_splits,min_samples_leaf=min_samples_leafs,max_features=max_featuress,criterion=criterions,random_state=RANDOM_STATE)
            clf1 = clf1.fit(X_train,y_train)
            pre=clf1.predict(X_test)
            acc=accuracy_score(y_test,pre)
            print(acc)
            print("training done")
            return acc

        if "svm-rbf" in algo_tune:
            var3=[]
            keys3=[]
            for key,value in li[u].items():
                if value !="":
                    var3.append(key)
                    keys3.append(value)
            if "C" in var3:
                indi11=var3.index('C')
                c=int(keys3[indi11])
            else:
                c=1.0
            if "gamma" in var3:
                indi12=var3.index('gamma')
                gammas=keys3[indi12]
            else:
                gammas='auto'
            clf1=svm.SVC(kernel='rbf', gamma=gammas, C=c,random_state=RANDOM_STATE)
            clf1 = clf1.fit(X_train,y_train)
            pre=clf1.predict(X_test)
            acc=accuracy_score(y_test,pre)
            print(acc)
            print("training done")
            return acc
        if "svm-linear" in algo_tune:
            var4=[]
            keys4=[]
            for key,value in li[u].items():
                if value !="":
                    var4.append(key)
                    keys4.append(value)
            if "C" in var4:
                indi13=var4.index('C')
                c=int(keys4[indi13])
            else:
                c=1.0
            if 'class_weight' in var4:
                indi14=var4.index('class_weight')
                class_weights=keys4[indi14]
            else:
                class_weights=None
            if "multi_class" in var4:
                indi15=var4.index('multi_class')
                multi_classs=keys4[indi15]
            else:
                multi_classs='ovr'
            clf1=svm.LinearSVC(C=c,class_weight=class_weights,multi_class=multi_classs,random_state=RANDOM_STATE)
            clf1 = clf1.fit(X_train,y_train)
            pre=clf1.predict(X_test)
            acc=accuracy_score(y_test,pre)
            print(acc)
            print("training done")
            return acc
        if "KNeighborsClassifier" in algo_tune:
            var5=[]
            keys5=[]
            for key,value in li[u].items():
                if value !="":
                    var5.append(key)
                    keys5.append(value)
            if 'n_neighbors' in var5:
                indi16=var5.index('n_neighbors')
                n_neighborss=int(keys5[indi16])
            else:
                n_neighborss=5
           # if 'class_weight' in var5:
           #     indi17=var5.index('class_weight')
           #     class_weights=keys5[indi17]
           # else:
           #     class_weights=None
            if 'algorithm' in var5:
                indi18=var5.index("algorithm")
                algorithms=keys5[indi18]
            else:
                algorithms='ovr'
            if 'p' in var5:
                indi19=var5.index('p')
                ps=int(keys5[indi19])
            else:
                ps=2
            clf1=KNeighborsClassifier(n_neighbors=n_neighborss,algorithm=algorithms,p=ps)
            clf1 = clf1.fit(X_train,y_train)
            pre=clf1.predict(X_test)
            acc=accuracy_score(y_test,pre)
            print(acc)
            print("training done")
            return acc
def model(algo,ex_dict,mid,algo_params,email):
    global X_train,X_test,y_train,y_test,filepath1,feature_count,sel_x,dummy,tt
    print(algo_params)
    li={}
    acc={}
    for i in algo_params:
        li[i]=algo_params[i]
        if "classification" in algo:
            acc[i]=algo_select(li)
            li.clear()
        else:
            acc[i]=algo_selectr(li)
            li.clear()
    print(acc)
    ss={"modelId":mid,"algo":algo,"accuracy":acc,"email":email,"status":"True"}
    success_json=json.dumps(ss)
    return success_json


'''
2)	Ler uma matriz A de duas dimensões
com 7 linhas e 7 colunas. Ao final
apresentar o total de elementos pares
existentes dentro da matriz.
'''
import random
matriz = [[0] * 7 for i in range(7)]

for lin in range(7):
    for col in range(7):
        matriz[lin][col] = random.randint(0,100)

conta = 0
for lin in range(7):
    for col in range(7):
        if matriz[lin][col] % 2 == 0:
            conta += 1

print(f"Numero de elementos: {conta}") #String Interpolation


"""Flask app for Cupcakes"""
from flask import Flask, render_template, jsonify, request
from flask_debugtoolbar import DebugToolbarExtension
from models import *

app = Flask(__name__)
app.app_context().push()
app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql:///blogly'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['SQLALCHEMY_ECHO'] = True
app.config['DEBUG_TB_INTERCEPT_REDIRECTS'] = False
app.config['SECRET_KEY'] = "sec_ret$0987"
debug = DebugToolbarExtension(app)

connect_db(app)

def serialize_cupcake(cupcake):
    """Serialize a cupcake SQLAlchemy obj to dictionary."""
    return {
        "id": cupcake.id,
        "flavor": cupcake.flavor,
        "size": cupcake.size,
        "rating": cupcake.rating,
        "image": cupcake.image
    }

@app.route("/")
def home_page():
    return render_template("homepage.html")

@app.route("/api/cupcakes")
def list_cupcakes():
    """Get data about all cupcakes"""
    cupcakes = Cupcake.query.all()
    serialized = [serialize_cupcake(c) for c in cupcakes]
    return jsonify(cupcakes=serialized)

@app.route("/api/cupcakes/<int:cupcake_id>")
def get_cupcake(cupcake_id):
    """Get data about a single cupcake"""
    cupcake = Cupcake.query.get_or_404(cupcake_id)
    serialized = serialize_cupcake(cupcake)
    return jsonify(cupcake=serialized)

@app.route("/api/cupcakes", methods=["POST"])
def create_cupcake():
    """Create a cupcake form data & return it"""
    flavor = request.json["flavor"]
    size = request.json["size"]
    rating = request.json["rating"]
    image = request.json["image"]
    new_cupcake = Cupcake(flavor=flavor, size=size, rating=rating, image=image)
    db.session.add(new_cupcake)
    db.session.commit()

    serialized = serialize_cupcake(new_cupcake)
    return(jsonify(cupcake=serialized), 201)

@app.route("/api/cupcakes/<int:cupcake_id>", methods=["PATCH"])
def update_cupcake(cupcake_id):
    """Update a cupcake"""
    cupcake = Cupcake.query.get_or_404(cupcake_id)
    cupcake.flavor = request.json["flavor"]
    cupcake.size = request.json["size"]
    cupcake.rating = request.json["rating"]
    cupcake.image = request.json["image"]
    db.session.commit()

    serialized = serialize_cupcake(cupcake)
    return jsonify(cupcake=serialized)

@app.route("/api/cupcakes/<int:cupcake_id>", methods=["DELETE"])
def delete_cupcake(cupcake_id):
    """Delete a cupcake"""
    cupcake = Cupcake.query.get_or_404(cupcake_id)
    db.session.delete(cupcake)
    db.session.commit()

    return jsonify(message="Deleted")


# USDA National Nutrient Database Integration
from __future__ import absolute_import, division, print_function
import requests

USDA_NDB_API_PATH = 'http://api.nal.usda.gov/ndb/'

def search(api_key, search_term):
    """ http://ndb.nal.usda.gov/ndb/doc/apilist/API-SEARCH.md

    Internal params of interest
    Sort the results by food name (n) or by search relevance (r)"""

    url_path = USDA_NDB_API_PATH + 'search/'
    payload = {'api_key': api_key, 'format' : 'json', 'q' : search_term, 'sort' : 'r', 'max' : '25','offset' : '0'}
    response = requests.get(url_path, params=payload)
    return response.json()

def food_report(api_key, ndbno):
    """ http://ndb.nal.usda.gov/ndb/doc/apilist/API-FOOD-REPORT.md

    Internal params of interest
    type    n   b (basic)   Report type: [b]asic or [f]ull or [s]tats"""

    url_path = USDA_NDB_API_PATH + 'reports/'
    payload = {'api_key': api_key, 'format' : 'json', 'ndbno' : ndbno, 'type' : 'b'}
    response = requests.get(url_path, params=payload)

    return response.json()

def kcal_for_ndbno(api_key, ndbno):
    """convenience function to get kcals for a food item"""

    NUTRIENT_ID_ENERGY = "208"
    report = food_report(api_key, ndbno)

    # error handling is for losers
    nutrients = report["report"]["food"]["nutrients"]
    for nutrient in nutrients:
        id = nutrient["nutrient_id"]
        if id == NUTRIENT_ID_ENERGY:
            return nutrient["value"]


from django.urls import path
from . import views

app_name = 'app'

urlpatterns = [
    path('register/', views.register, name='register'),
    path('login/', views.login, name='login'),
    path('logout/', views.logout, name='logout'),

    path('', views.home, name='home'),

    path('marketplace/', views.placements, name='placements'),
    path('marketplace/<placement_slug>/', views.placement_detail, name='placement-detail'),

    path('my-bids/', views.bid_summary, name='bid-summary' ),
    path('confirm-bids/', views.confirm_bids, name='confirm-bids'),

    path('dashboard/', views.dashboard, name='dashboard'),
    path('about/', views.about, name='about')
]


"""
Test conversions between image and dataframe.
"""
from itertools import product
import pytest
import numpy as np

from .._image import image_to_dataframe, dataframe_to_image

RND_SEED = 42

@pytest.mark.parametrize("number_of_bands", list(range(1, 20, 3)))
def test_image_to_dataframe(number_of_bands):
    """Test the conversion of image arrays to regular dataframes."""
    shape = (150, 200, number_of_bands)
    band_names = [f"a{i}" for i in range(number_of_bands)]

    X = np.random.default_rng(RND_SEED).random(shape)
    target = (
        np.random.default_rng(RND_SEED)
        .integers(0, 2, shape[0] * shape[1])
        .reshape(shape[0], shape[1])
    )

    for bands, y in product([band_names, None], [target, None]):
        df = image_to_dataframe(X=X, y=y, bands=bands)

        if y is None:
            assert df.shape == (shape[0] * shape[1], shape[2])
            for h, w in product([0, shape[0] - 1], [0, shape[1] - 1]):
                assert (df.loc[h, w] == X[h, w]).all()
        else:
            assert df.shape == (shape[0] * shape[1], shape[2] + 1)
            assert df.columns[-1] == "target"
            assert df["target"].dtype == int
            for h, w in product([0, shape[0] - 1], [0, shape[1] - 1]):
                assert (df.loc[h, w] == np.append(X[h, w], y[h, w])).all()

        if bands is None:
            assert df.columns[:number_of_bands].tolist() == list(range(number_of_bands))
        else:
            assert df.columns[:number_of_bands].tolist() == band_names

@pytest.mark.parametrize("number_of_bands", list(range(1, 20, 3)))
def test_dataframe_to_image(number_of_bands):
    shape = (150, 200, number_of_bands)
    band_names = [f"a{i}" for i in range(number_of_bands)]

    X = np.random.default_rng(RND_SEED).random(shape)
    target = (
        np.random.default_rng(RND_SEED)
        .integers(0, 2, shape[0] * shape[1])
        .reshape(shape[0], shape[1])
    )

    for bands, y in product([band_names, None], [target, None]):
        df = image_to_dataframe(X=X, y=y, bands=bands)
        target_ = "target" if y is not None else None
        X_, y_, bands_ = dataframe_to_image(df=df, target_feature=target_)

        if y is None:
            assert y_ is None
        else:
            assert type(y_) == np.ndarray
            assert (y_ == y).all()

        assert (bands_ == df.columns.drop(target_, errors="ignore")).all()
        assert type(X_) == np.ndarray
        assert (X_ == X).all()

        # Without "h", "w" features at all
        with pytest.raises(IndexError):
            dataframe_to_image(df=df.reset_index(drop=True), target_feature=target_)

        # With "h", "w" in columns
        with pytest.raises(IndexError):
            dataframe_to_image(df=df.reset_index(drop=True), target_feature=target_)

    # Test passing a subset of bands
    if len(band_names) > 2:
        df = image_to_dataframe(X=X, y=target, bands=band_names)
        bands_filter = np.random.choice(band_names, 2, replace=False)

        X_, y_, bands_ = dataframe_to_image(
            df=df, bands=bands_filter, target_feature="target"
        )

        assert (bands_ == bands_filter).all()
        assert type(y_) == np.ndarray
        assert type(X_) == np.ndarray
        assert (y_ == target).all()

        band_indices = [
            np.where(df.columns.drop("target") == b)[0][0] for b in bands_filter
        ]
        assert (X_ == X[:, :, band_indices]).all()


import yaml
import uuid
import time
import qrcode
from io import BytesIO
from flask import g
from app.models.coupon_finder_db_model import User, Category, Coupon
from sqlalchemy.orm import sessionmaker
from qiniu import Auth, put_file
from datetime import datetime
from . import coupon_finder_engine

with open('app/config.yml') as f:
    config = yaml.safe_load(f)

def get_coupon_category_id(keyword: str) -> int:
    card_package_id = get_db_session().query(Category).filter(Category.name == keyword).first().id
    return card_package_id

def get_card_package_id(open_id: str) -> int :
    card_package_id = get_db_session().query(User.card_package_id).filter(User.open_id == open_id).first()
    return int(card_package_id[0])

def get_user_id(open_id: str) -> int:
    user_id = get_db_session().query(User.id).filter(User.open_id == open_id).first()
    return int(user_id[0])

def get_released_coupon_merchant_id(coupon_id: int) -> int:
    '''
        发布优惠券的商家的用户ID
    '''
    merchant_id = get_db_session().query(Coupon.merchant_id).filter(Coupon.id == int(coupon_id)).first()[0]
    return int(merchant_id)

def generate_random_filename():
    random_uuid = uuid.uuid4()
    return str(random_uuid)

def upload_file(file_dir_path: str, file_name: str):
    q = Auth(config['qiniu']['accessKey'], config['qiniu']['secretKey'])
    #要上传的空间
    bucket_name = config['qiniu']['bucketName']
    #生成上传 Token，可以指定过期时间等
    token = q.upload_token(bucket_name, file_name, 3600)
    #要上传文件的本地路径
    put_file(token, file_name, file_dir_path + '/' + file_name, version='v2')

def get_current_ts():
    '''
    获取当前时间戳 单位：毫秒
    '''
    return int(time.time() * 1000)

def format_ts(ts: int):
    '''
        将时间戳格式化为字符串%Y-%m-%d %H:%M:%S
    '''
    # 假设给定的时间戳为1691143122611
    timestamp = ts / 1000  # 将毫秒转换为秒
    datetime_obj = datetime.fromtimestamp(timestamp)

    # 格式化为字符串
    formatted_time = datetime_obj.strftime('%Y-%m-%d %H:%M:%S')
    return formatted_time

def get_qrcode_byte_stream(data) -> bytes:
    qr = qrcode.QRCode(
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=10,
        border=4,
    )
    qr.add_data(data)
    qr.make(fit=True)

    img = qr.make_image(fill_color="black", back_color="white")

    # 将二维码图片以字节流的形式发送给前端
    img_byte_array = BytesIO()
    img.save(img_byte_array)
    img_byte_array.seek(0)
    return img_byte_array

def get_db_session():
    if 'db_session' not in g:
        g.db_session = sessionmaker(bind=coupon_finder_engine)()
    return g.db_session


import argparse
import io
import math
import random
import tarfile
import tempfile
from functools import partial, reduce
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

import psutil
import webdataset as wds

import numpy as np
import pandas as pd
import rioxarray
import xarray as xr
from deadtrees.utils.data_handling import make_blocks_vectorized, split_df
from PIL import Image
from tqdm.contrib.concurrent import process_map

random.seed(42)

Path.ls = lambda x: list(x.iterdir())

SHARDSIZE = 32
OVERSAMPLE_FACTOR = 2  # factor of random samples to dt + ndt samples

"""Summary:

This script builds up the final datasets for model training, validation and testing

The final datasets (combo) consists of three parts:
(1) tiles with classified deadtrees (frac > 0)
(2) tiles with non-deadtrees (aka healthy forest tiles, frac = 0)
(3) random other tiles (landuses: urban, arable, water etc., frac unknown, likely 0)

All shards are balanced to contain an equal amount of classified deadtree occurences to
allow a fair use of shards for traingin/ validation/ testing (composition of images should
be fair no matter the use)

Before the final dataset (combo) is created, various (temporary) datasets are created:
(1) train - raw, do not use
(2) train-balanced - balance the amounts of deadpixels (also filter to only include tiles with deadtrees)
(3) train-negsamples - a set of tiles that are guaranteed to contain non-dead trees
(4) train-randomsamples - a set of random other tiles

(X) combo: (2) + (4) [take turns]

"""

class Extractor:
    """Extract subtiles from rgbn or mask tile"""

    def __init__(self, *, tile_size: int = 256, source_dim: int = 2048):
        self.tile_size = tile_size
        self.source_dim = source_dim

    def __call__(self, t: Optional[xr.DataArray], *, n_bands: int):
        """Get data from tile, zeropad if necessary"""

        # default: all-zero in case no mask file exists

        if t is None:
            data = np.zeros((n_bands, self.source_dim, self.source_dim), dtype=np.uint8)
        else:
            data = np.zeros((n_bands, self.source_dim, self.source_dim), dtype=t.dtype)
            if (len(t.x) * len(t.y)) != (self.source_dim * self.source_dim):
                data[:, 0 : 0 + t.shape[1], 0 : 0 + t.shape[2]] = t.values
            else:
                data = t.values

        return make_blocks_vectorized(data, self.tile_size)

def _split_tile(
    image: Path,
    mask: Path,
    lu: Path,
    *,
    source_dim: int,
    tile_size: int,
    format: str,
    valid_subtiles: Optional[Iterable[str]] = None,
) -> List[Tuple[str, bytes, bytes]]:
    """Helper func for split_tiles"""

    extract = Extractor(tile_size=tile_size, source_dim=source_dim)

    n_bands = 4  # RGBN
    chunks = {"band": n_bands, "x": tile_size, "y": tile_size}
    with rioxarray.open_rasterio(image, chunks=chunks) as t:
        subtile_rgbn = extract(t, n_bands=n_bands)

    # process (optional) mask data
    if mask:
        chunks = {"band": n_bands, "x": tile_size, "y": tile_size}
        with rioxarray.open_rasterio(mask, chunks=chunks) as t:
            subtile_mask = extract(t, n_bands=1)
    else:
        subtile_mask = extract(None, n_bands=1)

    # process (optional) lu data
    if lu:
        chunks = {"band": n_bands, "x": tile_size, "y": tile_size}
        with rioxarray.open_rasterio(lu, chunks=chunks) as t:
            subtile_lu = extract(t, n_bands=1)
    else:
        # NOTE: convert to 1 for lu image (instead of 0 for masks)
        # TODO: check if this is the right thing to do in general
        subtile_lu = extract(None, n_bands=1) + 1

    samples = []
    if format == "TIFF":
        suffix = "tif"
    elif format == "PNG":
        suffix = "png"
    else:
        raise NotImplementedError

    for i in range(subtile_rgbn.shape[0]):
        subtile_name = f"{image.stem}_{i:03}"

        if np.min(subtile_rgbn[i]) != np.max(subtile_rgbn[i]):
            im = Image.fromarray(np.rollaxis(subtile_rgbn[i], 0, 3), "RGBA")
            im_mask = Image.fromarray(subtile_mask[i].squeeze())
            im_lu = Image.fromarray(subtile_lu[i].squeeze())

            im_byte_arr = io.BytesIO()
            im.save(im_byte_arr, format=format)
            im_byte_arr = im_byte_arr.getvalue()

            im_mask_byte_arr = io.BytesIO()
            im_mask.save(im_mask_byte_arr, format=format)
            im_mask_byte_arr = im_mask_byte_arr.getvalue()

            im_lu_byte_arr = io.BytesIO()
            im_lu.save(im_lu_byte_arr, format=format)
            im_lu_byte_arr = im_lu_byte_arr.getvalue()

            sample = {
                "__key__": subtile_name,
                f"rgbn.{suffix}": im_byte_arr,
                f"mask.{suffix}": im_mask_byte_arr,
                f"lu.{suffix}": im_lu_byte_arr,
                "txt": str(
                    round(
                        float(np.count_nonzero(subtile_mask[i]))
                        / (tile_size * tile_size)
                        * 100,
                        2,
                    )
                ),
            }
            if (valid_subtiles is None) or (subtile_name in valid_subtiles):
                samples.append(sample)
    return samples

def split_tiles(
    images, masks, lus, workers: int, shardpattern: str, **kwargs
) -> List[Any]:
    """Split tile into subtiles in parallel and save them to disk"""

    valid_subtiles = kwargs.get("valid_subtiles", None)

    stats = []
    with wds.ShardWriter(shardpattern, maxcount=SHARDSIZE) as sink:

        data = process_map(
            partial(_split_tile, **kwargs),
            images,
            masks,
            lus,
            max_workers=workers,
            chunksize=1,
        )

        for sample in reduce(lambda z, y: z + y, data):
            if sample:
                if valid_subtiles:
                    if sample["__key__"] in valid_subtiles:
                        sink.write(sample)
                        stats.append((sample["__key__"], sample["txt"], "1"))
                else:
                    if float(sample["txt"]) > 0:
                        sink.write(sample)
                        stats.append((sample["__key__"], sample["txt"], "1"))
                    else:
                        # not included in shard
                        stats.append((sample["__key__"], sample["txt"], "0"))

    return stats

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("image_dir", type=Path)
    parser.add_argument("mask_dir", type=Path)
    parser.add_argument("lu_dir", type=Path)
    parser.add_argument("outdir", type=Path)

    num_cores = psutil.cpu_count(logical=False)
    parser.add_argument(
        "--workers",
        dest="workers",
        type=int,
        default=num_cores,
        help="number of workers for parallel execution [def: %(default)s]",
    )

    parser.add_argument(
        "--source_dim",
        dest="source_dim",
        type=int,
        default=2048,
        help="size of input tiles [def: %(default)s]",
    )

    parser.add_argument(
        "--tile_size",
        dest="tile_size",
        type=int,
        default=256,
        help="size of final tiles that are then passed to the model [def: %(default)s]",
    )

    parser.add_argument(
        "--format",
        dest="format",
        type=str,
        default="TIFF",
        choices=["PNG", "TIFF"],
        help="target file format (PNG, TIFF) [def: %(default)s]",
    )

    parser.add_argument(
        "--tmp-dir",
        dest="tmp_dir",
        type=Path,
        default=None,
        help="use this location as tmp dir",
    )

    parser.add_argument(
        "--subdir",
        dest="sub_dir",
        default="train",
        help="use this location as sub_dir",
    )

    parser.add_argument(
        "--stats",
        dest="stats_file",
        type=Path,
        default=Path("stats.csv"),
        help="use this file to record stats",
    )

    args = parser.parse_args()

    args.outdir.mkdir(parents=True, exist_ok=True)
    Path(args.outdir / args.sub_dir).mkdir(parents=True, exist_ok=True)

    if args.tmp_dir:
        print(f"Using custom tmp dir: {args.tmp_dir}")
        Path(args.tmp_dir).mkdir(parents=True, exist_ok=True)

    if args.format == "TIFF":
        suffix = "tif"
    elif args.format == "PNG":
        suffix = "png"
    else:
        raise NotImplementedError

    SHUFFLE = True  # shuffle subtile order within shards (with fixed seed)

    # subtile_stats = split_tiles(train_files)
    images = sorted(args.image_dir.glob("*.tif"))
    masks = sorted(args.mask_dir.glob("*.tif"))
    lus = sorted(args.lu_dir.glob("*.tif"))

    image_names = {i.name for i in images}
    mask_names = {i.name for i in masks}
    lu_names = {i.name for i in lus}

    # limit set of images to images that have equivalent mask tiles
    train_images = [
        i
        for i in images
        if i.name in image_names.intersection(mask_names).intersection(lu_names)
    ]
    train_masks = [
        i
        for i in masks
        if i.name in mask_names.intersection(image_names).intersection(lu_names)
    ]
    train_lus = [
        i
        for i in lus
        if i.name in lu_names.intersection(mask_names).intersection(image_names)
    ]

    train_images = sorted(train_images)
    train_masks = sorted(train_masks)
    train_lus = sorted(train_lus)

    # print(len(train_images))
    # print(len(train_masks))
    # exit()
    # print(len(train_lus))

    cfg = dict(
        source_dim=args.source_dim,
        tile_size=args.tile_size,
        format=args.format,
    )

    subtile_stats = split_tiles(
        train_images,
        train_masks,
        train_lus,
        args.workers,
        str(args.outdir / args.sub_dir / "train-%06d.tar"),
        **cfg,
    )

    with open(args.outdir / args.stats_file, "w") as fout:
        fout.write("tile,frac,status\n")
        for i, (fname, frac, status) in enumerate(subtile_stats):
            line = f"{fname},{frac},{status}\n"
            fout.write(line)

    # rebalance shards so we get similar distributions in all shards
    with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmpdir:
        print(f"Created a temporary directory: {tmpdir}")

        print("Extract source tars")
        # untar input
        for tf_name in sorted((args.outdir / args.sub_dir).glob("train-00*.tar")):
            with tarfile.open(tf_name) as tf:
                tf.extractall(tmpdir)

        print("Write balanced shards from deadtree samples")
        df = pd.read_csv(args.outdir / args.stats_file)

        df = df[df.status > 0]

        n_valid = len(df)
        splits = split_df(df, SHARDSIZE)

        # preserve last shard if more than 50% of values are present
        if SHARDSIZE // 2 < len(splits[-1]) < SHARDSIZE:
            # fill last shard with duplicates (not ideal...)
            n_missing = SHARDSIZE - len(splits[-1])
            # df_extra = splits[-1].sample(n=n_missing, random_state=42)
            splits[-1].extend(np.random.choice(splits[-1], size=n_missing).tolist())

        # drop incomplete shards
        splits = [x for x in splits if len(x) == SHARDSIZE]
        assert len(splits) > 0, "Something went wrong"

        for s_cnt, s in enumerate(splits):

            with tarfile.open(
                args.outdir / args.sub_dir / f"train-balanced-{s_cnt:06}.tar", "w"
            ) as dst:

                if SHUFFLE:
                    random.shuffle(s)
                for i in s:
                    dst.add(f"{tmpdir}/{i}.mask.{suffix}", f"{i}.mask.{suffix}")
                    dst.add(f"{tmpdir}/{i}.lu.{suffix}", f"{i}.lu.{suffix}")
                    dst.add(f"{tmpdir}/{i}.rgbn.{suffix}", f"{i}.rgbn.{suffix}")
                    dst.add(f"{tmpdir}/{i}.txt", f"{i}.txt")

    # create sets for random tile dataset
    # use all subtiles not covered in train

    n_subtiles = (args.source_dim // args.tile_size) ** 2

    all_subtiles = []
    for image_name in image_names:
        all_subtiles.extend(
            [f"{Path(image_name).stem}_{c:03}" for c in range(n_subtiles)]
        )
    all_subtiles = set(all_subtiles)

    n_samples = n_valid * OVERSAMPLE_FACTOR
    random_subtiles = random.sample(
        tuple(all_subtiles - set([x[0] for x in subtile_stats if int(x[2]) == 1])),
        n_samples,
    )

    # the necessary tile to process
    random_tiles = sorted(list(set([x[:-4] for x in random_subtiles])))

    all_images = sorted(args.image_dir.glob("*.tif"))
    random_images = [x for x in all_images if x.stem in random_tiles]

    print("STATS")
    print(len(all_subtiles))
    print(len(subtile_stats))
    print(len(random_subtiles))
    print(len(random_images))

    cfg = dict(
        source_dim=args.source_dim,
        tile_size=args.tile_size,
        format=args.format,
        valid_subtiles=random_subtiles,  # subset data with random selection of subtiles
    )

    random_images_names = {i.name for i in random_images}
    random_lus = [i for i in lus if i.name in random_images_names]

    subtile_stats_rnd = split_tiles(
        random_images,
        [None] * len(random_images),
        random_lus,
        args.workers,
        str(args.outdir / args.sub_dir / "train-randomsamples-%06d.tar"),
        **cfg,
    )

    stats_file_rnd = Path(args.stats_file.stem + "_rnd.csv")
    with open(args.outdir / stats_file_rnd, "w") as fout:
        fout.write("tile,frac,status\n")
        for i, (fname, frac, status) in enumerate(subtile_stats_rnd):
            line = f"{fname},{frac},{status}\n"
            fout.write(line)

    # also create combo dataset
    # source A: train-balanced, source B: randomsample
    # NOTE: combo dataset has double the default shardsize (2*128), samples alternate between regular and random sample
    train_balanced_shards = [
        str(x) for x in sorted((args.outdir / args.sub_dir).glob("train-balanced*"))
    ]
    train_balanced_shards_rnd = [
        str(x) for x in sorted((args.outdir / args.sub_dir).glob("train-random*"))
    ]
    train_balanced_shards_rnd = train_balanced_shards_rnd[: len(train_balanced_shards)]

    shardpattern = str(args.outdir / args.sub_dir / "train-combo-%06d.tar")

    with wds.ShardWriter(shardpattern, maxcount=SHARDSIZE * 2) as sink:
        for shardA, shardB in zip(train_balanced_shards, train_balanced_shards_rnd):

            for sA, sB in zip(wds.WebDataset(shardA), wds.WebDataset(shardB)):
                sink.write(sA)
                sink.write(sB)

    # remove everything but train & combo
    for filename in (args.outdir / args.sub_dir).glob("train-random*"):
        filename.unlink()
    for filename in (args.outdir / args.sub_dir).glob("train-balanced*"):
        filename.unlink()
    for filename in (args.outdir / args.sub_dir).glob("train-0*"):
        filename.unlink()

if __name__ == "__main__":
    main()