# -*- coding: utf-8 -*-
"""
Class definition of YOLO_v3 style detection model on image and video
"""
import argparse
import colorsys
import os
from timeit import default_timer as timer

import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw, ImageTk

from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body
from yolo3.utils import letterbox_image
import os
from keras.utils import multi_gpu_model

import blue_1
import day10
import tkinter as tk
from tkinter import filedialog

#修改：用来记录车辆位置并记录的变量

left_re=0
top_re=0
flag_num=0
bottom_re=0
right_re=0
flag_count=0
count_num=0
flag_count1=0
count_num1=0
over_speed_num = 0
car_num = 0
car_counttime = 0

license_num = ["","","","",""]

class YOLO(object):
    _defaults = {
        "model_path": 'model_data/yolo.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/coco_classes.txt',
        "score" : 0.3,
        "iou" : 0.45,
        "model_image_size" : (416, 416),
        "gpu_num" : 1,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults) # set up default values
        self.__dict__.update(kwargs) # and update with user overrides
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.sess = K.get_session()
        self.boxes, self.scores, self.classes = self.generate()

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def generate(self):
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)
        is_tiny_version = num_anchors==6 # default setting
        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = tiny_yolo_body(Input(shape=(None,None,3)), num_anchors//2, num_classes) \
                if is_tiny_version else yolo_body(Input(shape=(None,None,3)), num_anchors//3, num_classes)
            self.yolo_model.load_weights(self.model_path) # make sure model, anchors and classes match
        else:
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                num_anchors/len(self.yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # Generate colors for drawing bounding boxes.
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        np.random.seed(None)  # Reset seed to default.

        # Generate output tensor targets for filtered bounding boxes.
        self.input_image_shape = K.placeholder(shape=(2, ))
        if self.gpu_num>=2:
            self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num)
        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                len(self.class_names), self.input_image_shape,
                score_threshold=self.score, iou_threshold=self.iou)
        return boxes, scores, classes

    def detect_image(self, image):
        start = timer()

        if self.model_image_size != (None, None):
            assert self.model_image_size[0]%32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1]%32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        print(image_data.shape)
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                K.learning_phase(): 0
            })

        print('Found {} boxes for {}'.format(len(out_boxes), 'img'))

        font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness = (image.size[0] + image.size[1]) // 300

        for i, c in reversed(list(enumerate(out_classes))):
            predicted_class = self.class_names[c]
            box = out_boxes[i]
            score = out_scores[i]

            label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)
            label1 = label[0:3]
            if label1 != "tra" and label1 != "per" and label1!="umb" and label1!="bic" and label1!="tru" and \
                    label1!="han" and label1!="fir" and label1!="pot" and label1!="bus":
                top, left, bottom, right = box
                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
                right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
               # print(label, (left, top), (right, bottom))

                # 修改：用来记录车辆位置并记录
                global left_re

                global top_re

                global flag_num

                global bottom_re

                global right_re

                global car_counttime

                global car_num

                if (left_re != 0 and top_re != 0) and (
                        (top_re - top > 20 and top_re - top < 40) or (top_re - top < -20 and top_re - top > -40)) and (
                        left<720 and left>520) and(
                        top > 600 and top < 675) and flag_num < 1:

                    top_re = top
                    right_re=right
                    left_re=left
                    bottom_re=bottom
                    flag_num = flag_num + 1

                if((left<578  and left>520) and (top>600 and top<675)) and flag_num==0:
                    left_re = left
                    top_re = top

                if left > 520 and (
                        top > 600 and top < 640) and car_counttime<=1:
                    car_num = car_num +1
                    car_counttime = 1

                if top - label_size[1] >= 0:
                    text_origin = np.array([left, top - label_size[1]])
                else:
                    text_origin = np.array([left, top + 1])

                # My kingdom for a good redistributable image drawing library.
                for i in range(thickness):
                    draw.rectangle(
                        [left + i, top + i, right - i, bottom - i],
                        outline=self.colors[c])
                draw.rectangle(
                    [tuple(text_origin), tuple(text_origin + label_size)],
                    fill=self.colors[c])
                draw.text(text_origin, label, fill=(0, 0, 0), font=font)
                del draw
        print("**************"+str(car_num))
        if car_counttime==1:
            car_counttime = car_counttime +2
        if car_counttime>1:
            car_counttime = car_counttime + 1
        if car_counttime>=11:
            car_counttime=0
        end = timer()
        print(end - start)
        return image

    def close_session(self):
        self.sess.close()

def detect_video(yolo, video_path, output_path=""):
    import cv2
    # 更改视频2
    window = tk.Tk()
    window.title('正在识别中')
    window.geometry('1900x750')
    canvas = tk.Canvas(window, width=1400, height=788, bg="white")
    canvas.place(x=20, y=90)
    canvas1 = tk.Canvas(window, width=600, height=16, bg="white")
    canvas1.place(x=2, y=2)

    vid = cv2.VideoCapture(video_path)
    #vid = cv2.VideoCapture("video-02.mp4")
    if not vid.isOpened():
        raise IOError("Couldn't open webcam or video")
    video_FourCC    = int(vid.get(cv2.CAP_PROP_FOURCC))
    video_fps       = vid.get(cv2.CAP_PROP_FPS)
    fps_total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
    fps_pos =  0
    video_size      = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                        int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    output_path="video-01.mp4";
    isOutput = True if output_path != "" else False
    if isOutput:
        print("!!! TYPE:", type(output_path), type(video_FourCC), type(video_fps), type(video_size))
        out = cv2.VideoWriter("video-01.mp4", video_FourCC, video_fps, video_size)
    accum_time = 0
    curr_fps = 0
    fps = "FPS: ??"
    prev_time = timer()
    while True:
        return_value, frame = vid.read()
        cv2.rectangle(frame, (520, 600), (720,675), (255, 0, 0))
        cv2.putText(frame, "check_area", (520, 600), cv2.FONT_HERSHEY_COMPLEX, 2, (255, 0, 0), 2, 8, 0)
        image = Image.fromarray(frame)

        image = yolo.detect_image(image)
        result = np.asarray(image)
        curr_time = timer()
        exec_time = curr_time - prev_time
        prev_time = curr_time
        accum_time = accum_time + exec_time
        curr_fps = curr_fps + 1

        #修改截取保存图片1
        global left_re

        global top_re

        global flag_num

        global bottom_re

        global right_re

        global flag_count

        global count_num

        global flag_count1

        global count_num1

        global license_num

        global over_speed_num

        # 更改视频1

        pilImage = image
        pilImage = pilImage.resize((1500, 600), Image.ANTIALIAS)

        tkImage = ImageTk.PhotoImage(image=pilImage)
        canvas.create_image(0, 0, anchor='nw', image=tkImage)
        window.update_idletasks()
        window.update()

        if accum_time > 1:
            accum_time = accum_time - 1
            fps = "FPS: " + str(curr_fps)
            curr_fps = 0
        cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.50, color=(255, 0, 0), thickness=2)
        #cv2.namedWindow("result", cv2.WINDOW_NORMAL)
        #cv2.imshow("result", result)
        # 修改截取保存图片1
        if flag_num == 1:
            print("top_re:", top_re, "bottom_re", bottom_re, "left_re:", left_re, "right_re:", right_re)
            cutting = result[top_re:bottom_re, left_re:right_re]
            cv2.imwrite('cut'+str(over_speed_num)+'.jpg', cutting)
            flag_num = 10
            flag_count = 1
        if flag_count==1:
            count_num = count_num + 1
        if count_num==5:
            blue_1.check_license_plate(over_speed_num)
            flag_count1 = 1
        if flag_count1==1:
            count_num1 = count_num1+1
        if count_num1 == 5:
            license_num[over_speed_num] = day10.distinguish_license_plate()
            left_re = 0
            top_re = 0
            flag_num = 0
            bottom_re = 0
            right_re = 0
            flag_count = 0
            count_num = 0
            flag_count1 = 0
            count_num1 = 0
            over_speed_num = over_speed_num + 1

        if isOutput:
            out.write(result)
            fps_pos = fps_pos + 1
            fill_line = canvas1.create_rectangle(1.5, 1.5, 0, 23, width=0, fill="green")
            raise_data = 600 * (fps_pos / fps_total)
            canvas1.coords(fill_line, (0, 0, raise_data, 60))
            window.update()
            if fps_pos == fps_total - 1:
                window.destroy()
                break
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    yolo.close_session()

def dis_finish():
    import cv2
    # 更改视频2
    window = tk.Tk()
    window.title('成果展示')
    window.geometry('1900x750')
    canvas = tk.Canvas(window, width=1000, height=562, bg="white")
    canvas.place(x=20, y=20)
    global license_num
    global over_speed_num
    global car_num
    i=0
    label_text = tk.Label(window, text="下面车道总流量："+str(car_num))
    label_text.place(x=1020, y=5)
    while True:
        img = Image.open('cut'+str(i)+'.jpg')  # 打开图片
        photo = ImageTk.PhotoImage(img)
        label_img = tk.Label(window, image=photo)
        label_img.place(x=1020, y=155+500*i)
        license_num[i] = '超速车牌号：' + license_num[i]
        label_text = tk.Label(window, text=license_num[i].split('.', 1)[0])
        label_text.place(x=1020, y=55+500*i)
        label_text1 = tk.Label(window, text=license_num[i].split('.', 1)[1])
        label_text1.place(x=1020, y=105+500*i)
        i = i+1
        if i>over_speed_num-1:
            break
    vid = cv2.VideoCapture("video-01.mp4")
    if not vid.isOpened():
        raise IOError("Couldn't open webcam or video")
    accum_time = 0
    curr_fps = 0
    fps = "FPS: ??"
    prev_time = timer()
    while True:
        return_value, frame = vid.read()
        image = Image.fromarray(frame)
        curr_time = timer()
        exec_time = curr_time - prev_time
        prev_time = curr_time
        accum_time = accum_time + exec_time
        curr_fps = curr_fps + 1

        # 更改视频1

        pilImage = image
        pilImage = pilImage.resize((1000, 562), Image.ANTIALIAS)

        tkImage = ImageTk.PhotoImage(image=pilImage)
        canvas.create_image(0, 0, anchor='nw', image=tkImage)
        window.update_idletasks()
        window.update()
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

def hit_start():
    window.destroy()
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
    '''
    Command line options
    '''
    parser.add_argument(
        "--input", nargs='?', type=str, required=False, default='./path2your_video',
        help="Video input path"
    )

    parser.add_argument(
        "--output", nargs='?', type=str, default="",
        help="[Optional] Video output path"
    )
    FLAGS = parser.parse_args()

    detect_video(YOLO(**vars(FLAGS)), "video-02.mp4", '')
    dis_finish()
def hit_start_1():
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
    '''
    Command line options
    '''
    parser.add_argument(
        "--input", nargs='?', type=str, required=False, default='./path2your_video',
        help="Video input path"
    )

    parser.add_argument(
        "--output", nargs='?', type=str, default="",
        help="[Optional] Video output path"
    )
    FLAGS = parser.parse_args()
    f = tk.filedialog.askopenfilename()
    for i in range(len(f)):
        if f[i]=='/':
            f = f[:i]+'\\'+f[i+1:]
    print(f)
    window.destroy()
    detect_video(YOLO(**vars(FLAGS)), f, '')

    dis_finish()
if __name__ == '__main__':

    window = tk.Tk()
    window.title('车辆检测')
    window.geometry('500x300')
    b = tk.Button(window, text='固有视频测试', font=('Arial', 12), width=10, height=1, command=hit_start)
    b.pack()
    b1 = tk.Button(window, text='本地视频检测', font=('Arial', 12), width=10, height=1, command=hit_start_1)
    b1.pack()
    window.mainloop()


from models import unet, segnet
import cv2
import numpy as np
import argparse
from utils.segdata_generator import generator

def predict_segmentation():
    n_classes = 2
    images_path = '/home/deep/datasets/'
    val_file = './data/seg_test.txt'
    input_height = 256
    input_width = 256

    if args.model == 'unet':
        m = unet.Unet(n_classes, input_height=input_height, input_width=input_width)
    elif args.model == 'segnet':
        m = segnet.SegNet(n_classes, input_height=input_height, input_width=input_width)
    else:
        raise ValueError('Do not support {}'.format(args.model))

    m.load_weights("./results/{}_weights.h5".format(args.model))
    m.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

    colors = np.array([[0, 0, 0], [255, 255, 255]])
    i = 0
    for x, y in generator(images_path, val_file, 1, n_classes, input_height, input_width):
        pr = m.predict(x)[0]
        pr = pr.reshape((input_height, input_width, n_classes)).argmax(axis=2)
        seg_img = np.zeros((input_height, input_width, 3))
        for c in range(n_classes):
            seg_img[:, :, 0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8')
            seg_img[:, :, 1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8')
            seg_img[:, :, 2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8')
        cv2.imshow('test', seg_img)
        cv2.imwrite('./output/{}.jpg'.format(i), seg_img)
        i += 1
        cv2.waitKey(30)

if __name__ == '__main__':
    parse = argparse.ArgumentParser(description='command for training segmentation models with keras')
    parse.add_argument('--model', type=str, default='unet', help='support unet, segnet')
    args = parse.parse_args()
    predict_segmentation()


from __future__ import absolute_import, division, print_function

import logging
import numbers
from collections import defaultdict
from itertools import repeat
from multiprocessing import cpu_count
from operator import getitem

import dask
import numpy as np
import packaging.version
from dask.base import tokenize
from dask.delayed import delayed
from dask.distributed import as_completed
from dask.utils import derived_from
from sklearn import model_selection
from sklearn.base import BaseEstimator, MetaEstimatorMixin, clone, is_classifier
from sklearn.exceptions import NotFittedError
from sklearn.model_selection._search import BaseSearchCV, _check_param_grid
from sklearn.model_selection._split import (
    BaseShuffleSplit,
    KFold,
    LeaveOneGroupOut,
    LeaveOneOut,
    LeavePGroupsOut,
    LeavePOut,
    PredefinedSplit,
    StratifiedKFold,
    _BaseKFold,
    _CVIterableWrapper,
)
from sklearn.pipeline import FeatureUnion
try:
    from imblearn.pipeline import Pipeline
except:
    from sklearn.pipeline import Pipeline
from sklearn.utils.metaestimators import if_delegate_has_method
from sklearn.utils.multiclass import type_of_target
from sklearn.utils.validation import _num_samples

from .._compat import SK_VERSION, check_is_fitted
from ._normalize import normalize_estimator
from .methods import (
    MISSING,
    create_cv_results,
    cv_extract,
    cv_extract_params,
    cv_n_samples,
    cv_split,
    feature_union,
    feature_union_concat,
    feature_union_y,
    fit,
    fit_and_score,
    fit_best,
    fit_transform,
    pipeline,
    score,
)
from .utils import DeprecationDict, is_dask_collection, to_indexable, to_keys, unzip

logger = logging.getLogger(__name__)

try:
    from cytoolz import get, pluck
except ImportError:  # pragma: no cover
    from toolz import get, pluck

__all__ = ["GridSearchCV", "RandomizedSearchCV"]

if SK_VERSION <= packaging.version.parse("0.21.dev0"):

    _RETURN_TRAIN_SCORE_DEFAULT = "warn"

    def handle_deprecated_train_score(results, return_train_score):
        if return_train_score == "warn":
            results = DeprecationDict(results)
            message = (
                "You are accessing a training score ({!r}), "
                "which will not be available by default any more in "
                "sklearn 0.21. If you need training scores, please "
                "set return_train_score=True"
            )
            for key in results:
                if key.endswith("_train_score"):
                    results.add_warning(key, message.format(key), FutureWarning)
        return results

else:
    _RETURN_TRAIN_SCORE_DEFAULT = False

    def handle_deprecated_train_score(results, return_train_score):
        return results

class TokenIterator:
    def __init__(self, base_token):
        self.token = base_token
        self.counts = defaultdict(int)

    def __call__(self, est):
        typ = type(est)
        c = self.counts[typ]
        self.counts[typ] += 1
        return self.token if c == 0 else self.token + str(c)

def map_fit_params(dsk, fit_params):
    if fit_params:
        # A mapping of {name: (name, graph-key)}
        param_values = to_indexable(*fit_params.values(), allow_scalars=True)
        fit_params = {
            k: (k, v) for (k, v) in zip(fit_params, to_keys(dsk, *param_values))
        }
    else:
        fit_params = {}

    return fit_params

def build_graph(
    estimator,
    cv,
    scorer,
    candidate_params,
    X,
    y=None,
    groups=None,
    fit_params=None,
    iid=True,
    refit=True,
    error_score="raise",
    return_train_score=_RETURN_TRAIN_SCORE_DEFAULT,
    cache_cv=True,
    multimetric=False,
):
    # This is provided for compatibility with TPOT. Remove
    # once TPOT is updated and requires a dask-ml>=0.13.0
    def decompress_params(fields, params):
        return [{k: v for k, v in zip(fields, p) if v is not MISSING} for p in params]

    fields, tokens, params = normalize_params(candidate_params)
    dsk, keys, n_splits, main_token = build_cv_graph(
        estimator,
        cv,
        scorer,
        candidate_params,
        X,
        y=y,
        groups=groups,
        fit_params=fit_params,
        iid=iid,
        error_score=error_score,
        return_train_score=return_train_score,
        cache_cv=cache_cv,
    )
    cv_name = "cv-split-" + main_token
    if iid:
        weights = "cv-n-samples-" + main_token
        dsk[weights] = (cv_n_samples, cv_name)
        scores = keys[1:]
    else:
        scores = keys

    cv_results = "cv-results-" + main_token
    candidate_params_name = "cv-parameters-" + main_token
    dsk[candidate_params_name] = (decompress_params, fields, params)
    if multimetric:
        metrics = list(scorer.keys())
    else:
        metrics = None
    dsk[cv_results] = (
        create_cv_results,
        scores,
        candidate_params_name,
        n_splits,
        error_score,
        weights,
        metrics,
    )
    keys = [cv_results]
    return dsk, keys, n_splits

def build_cv_graph(
    estimator,
    cv,
    scorer,
    candidate_params,
    X,
    y=None,
    groups=None,
    fit_params=None,
    iid=True,
    error_score="raise",
    return_train_score=_RETURN_TRAIN_SCORE_DEFAULT,
    cache_cv=True,
):
    X, y, groups = to_indexable(X, y, groups)
    cv = check_cv(cv, y, is_classifier(estimator))
    # "pairwise" estimators require a different graph for CV splitting
    is_pairwise = getattr(estimator, "_pairwise", False)

    dsk = {}
    X_name, y_name, groups_name = to_keys(dsk, X, y, groups)
    n_splits = compute_n_splits(cv, X, y, groups)

    fit_params = map_fit_params(dsk, fit_params)

    fields, tokens, params = normalize_params(candidate_params)
    main_token = tokenize(
        normalize_estimator(estimator),
        fields,
        params,
        X_name,
        y_name,
        groups_name,
        fit_params,
        cv,
        error_score == "raise",
        return_train_score,
    )

    cv_name = "cv-split-" + main_token
    dsk[cv_name] = (cv_split, cv, X_name, y_name, groups_name, is_pairwise, cache_cv)

    if iid:
        weights = "cv-n-samples-" + main_token
        dsk[weights] = (cv_n_samples, cv_name)
    else:
        weights = None

    scores = do_fit_and_score(
        dsk,
        main_token,
        estimator,
        cv_name,
        fields,
        tokens,
        params,
        X_name,
        y_name,
        fit_params,
        n_splits,
        error_score,
        scorer,
        return_train_score,
    )
    keys = [weights] + scores if weights else scores
    return dsk, keys, n_splits, main_token

def build_refit_graph(estimator, X, y, best_params, fit_params):
    X, y = to_indexable(X, y)
    dsk = {}
    X_name, y_name = to_keys(dsk, X, y)

    fit_params = map_fit_params(dsk, fit_params)
    main_token = tokenize(normalize_estimator(estimator), X_name, y_name, fit_params)

    best_estimator = "best-estimator-" + main_token
    if fit_params:
        fit_params = (
            dict,
            (zip, list(fit_params.keys()), list(pluck(1, fit_params.values()))),
        )
    dsk[best_estimator] = (
        fit_best,
        clone(estimator),
        best_params,
        X_name,
        y_name,
        fit_params,
    )
    return dsk, [best_estimator]

def normalize_params(params):
    """Take a list of dictionaries, and tokenize/normalize."""
    # Collect a set of all fields
    fields = set()
    for p in params:
        fields.update(p)
    fields = sorted(fields)

    params2 = list(pluck(fields, params, MISSING))
    # Non-basic types (including MISSING) are unique to their id
    tokens = [
        tuple(x if isinstance(x, (int, float, str)) else id(x) for x in p)
        for p in params2
    ]

    return fields, tokens, params2

def _get_fit_params(cv, fit_params, n_splits):
    if not fit_params:
        return [(n, None) for n in range(n_splits)]
    keys = []
    vals = []
    for name, (full_name, val) in fit_params.items():
        vals.append(val)
        keys.append((name, full_name))
    return [(n, (cv_extract_params, cv, keys, vals, n)) for n in range(n_splits)]

def _group_fit_params(steps, fit_params):
    param_lk = {n: {} for n, _ in steps}
    for pname, pval in fit_params.items():
        step, param = pname.split("__", 1)
        param_lk[step][param] = pval
    return param_lk

def do_fit_and_score(
    dsk,
    main_token,
    est,
    cv,
    fields,
    tokens,
    params,
    X,
    y,
    fit_params,
    n_splits,
    error_score,
    scorer,
    return_train_score,
):
    if not isinstance(est, Pipeline):
        # Fitting and scoring can all be done as a single task
        n_and_fit_params = _get_fit_params(cv, fit_params, n_splits)

        est_type = type(est).__name__.lower()
        est_name = "%s-%s" % (est_type, main_token)
        score_name = "%s-fit-score-%s" % (est_type, main_token)
        dsk[est_name] = est

        seen = {}
        m = 0
        out = []
        out_append = out.append

        for t, p in zip(tokens, params):
            if t in seen:
                out_append(seen[t])
            else:
                for n, fit_params in n_and_fit_params:
                    dsk[(score_name, m, n)] = (
                        fit_and_score,
                        est_name,
                        cv,
                        X,
                        y,
                        n,
                        scorer,
                        error_score,
                        fields,
                        p,
                        fit_params,
                        return_train_score,
                    )
                seen[t] = (score_name, m)
                out_append((score_name, m))
                m += 1
        scores = [k + (n,) for n in range(n_splits) for k in out]
    else:
        X_train = (cv_extract, cv, X, y, True, True)
        X_test = (cv_extract, cv, X, y, True, False)
        y_train = (cv_extract, cv, X, y, False, True)
        y_test = (cv_extract, cv, X, y, False, False)

        # Fit the estimator on the training data
        X_trains = [X_train] * len(params)
        y_trains = [y_train] * len(params)
        fit_ests = do_fit(
            dsk,
            TokenIterator(main_token),
            est,
            cv,
            fields,
            tokens,
            params,
            X_trains,
            y_trains,
            fit_params,
            n_splits,
            error_score,
        )

        score_name = "score-" + main_token

        scores = []
        scores_append = scores.append
        for n in range(n_splits):
            if return_train_score:
                xtrain = X_train + (n,)
                ytrain = y_train + (n,)
            else:
                xtrain = ytrain = None

            xtest = X_test + (n,)
            ytest = y_test + (n,)

            for (name, m) in fit_ests:
                dsk[(score_name, m, n)] = (
                    score,
                    (name, m, n),
                    xtest,
                    ytest,
                    xtrain,
                    ytrain,
                    scorer,
                    error_score,
                )
                scores_append((score_name, m, n))
    return scores

def do_fit(
    dsk,
    next_token,
    est,
    cv,
    fields,
    tokens,
    params,
    Xs,
    ys,
    fit_params,
    n_splits,
    error_score,
):
    if issubclass(type(est), Pipeline) and params is not None:
        return _do_pipeline(
            dsk,
            next_token,
            est,
            cv,
            fields,
            tokens,
            params,
            Xs,
            ys,
            fit_params,
            n_splits,
            error_score,
            False,
        )
    else:
        n_and_fit_params = _get_fit_params(cv, fit_params, n_splits)

        if params is None:
            params = tokens = repeat(None)
            fields = None

        token = next_token(est)
        est_type = type(est).__name__.lower()
        est_name = "%s-%s" % (est_type, token)
        fit_name = "%s-fit-%s" % (est_type, token)
        dsk[est_name] = est

        seen = {}
        m = 0
        out = []
        out_append = out.append

        for X, y, t, p in zip(Xs, ys, tokens, params):
            if (X, y, t) in seen:
                out_append(seen[X, y, t])
            else:
                for n, fit_params in n_and_fit_params:
                    dsk[(fit_name, m, n)] = (
                        fit,
                        est_name,
                        X + (n,),
                        y + (n,),
                        error_score,
                        fields,
                        p,
                        fit_params,
                    )
                seen[(X, y, t)] = (fit_name, m)
                out_append((fit_name, m))
                m += 1

        return out

def do_fit_transform(
    dsk,
    next_token,
    est,
    cv,
    fields,
    tokens,
    params,
    Xs,
    ys,
    fit_params,
    n_splits,
    error_score,
):
    if issubclass(type(est), Pipeline) and params is not None:
        return _do_pipeline(
            dsk,
            next_token,
            est,
            cv,
            fields,
            tokens,
            params,
            Xs,
            ys,
            fit_params,
            n_splits,
            error_score,
            True,
        )
    elif issubclass(type(est), FeatureUnion) and params is not None:
        return _do_featureunion(
            dsk,
            next_token,
            est,
            cv,
            fields,
            tokens,
            params,
            Xs,
            ys,
            fit_params,
            n_splits,
            error_score,
        )
    else:
        n_and_fit_params = _get_fit_params(cv, fit_params, n_splits)

        if params is None:
            params = tokens = repeat(None)
            fields = None

        name = type(est).__name__.lower()
        token = next_token(est)
        fit_Xt_name = "%s-fit-transform-%s" % (name, token)
        fit_name = "%s-fit-%s" % (name, token)
        Xt_name = "%s-transform-%s" % (name, token)
        yt_name = "%s-ytransform-%s" % (name, token)
        est_name = "%s-%s" % (type(est).__name__.lower(), token)
        dsk[est_name] = est

        seen = {}
        m = 0
        out = []
        out_append = out.append

        for X, y, t, p in zip(Xs, ys, tokens, params):
            if (X, y, t) in seen:
                out_append(seen[X, y, t])
            else:
                for n, fit_params in n_and_fit_params:
                    dsk[(fit_Xt_name, m, n)] = (
                        fit_transform,
                        est_name,
                        X + (n,),
                        y + (n,),
                        error_score,
                        fields,
                        p,
                        fit_params,
                    )
                    dsk[(fit_name, m, n)] = (getitem, (fit_Xt_name, m, n), 0)
                    dsk[(Xt_name, m, n)] = (getitem, (fit_Xt_name, m, n), 1)
                    dsk[(yt_name, m, n)] = (getitem, (fit_Xt_name, m, n), 2)
                seen[X, y, t] = m
                out_append(m)
                m += 1

        return [(fit_name, i) for i in out], [(Xt_name, i) for i in out], [(yt_name, i) for i in out]

def _group_subparams(steps, fields, ignore=()):
    # Group the fields into a mapping of {stepname: [(newname, orig_index)]}
    field_to_index = dict(zip(fields, range(len(fields))))
    step_fields_lk = {s: [] for s, _ in steps}
    for f in fields:
        if "__" in f:
            step, param = f.split("__", 1)
            if step in step_fields_lk:
                step_fields_lk[step].append((param, field_to_index[f]))
                continue
        if f not in step_fields_lk and f not in ignore:
            raise ValueError("Unknown parameter: `%s`" % f)
    return field_to_index, step_fields_lk

def _group_ids_by_index(index, tokens):
    id_groups = []

    def new_group():
        o = []
        id_groups.append(o)
        return o.append

    _id_groups = defaultdict(new_group)
    for n, t in enumerate(pluck(index, tokens)):
        _id_groups[t](n)
    return id_groups

def _do_fit_step(
    dsk,
    next_token,
    step,
    cv,
    fields,
    tokens,
    params,
    Xs,
    ys,
    fit_params,
    n_splits,
    error_score,
    step_fields_lk,
    fit_params_lk,
    field_to_index,
    step_name,
    none_passthrough,
    is_transform,
):
    sub_fields, sub_inds = map(list, unzip(step_fields_lk[step_name], 2))
    sub_fit_params = fit_params_lk[step_name]

    if step_name in field_to_index:
        # The estimator may change each call
        new_fits = {}
        new_Xs = {}
        new_ys = {}
        est_index = field_to_index[step_name]

        for ids in _group_ids_by_index(est_index, tokens):
            # Get the estimator for this subgroup
            sub_est = params[ids[0]][est_index]
            if sub_est is MISSING:
                sub_est = step

            # If an estimator is `None`, there's nothing to do
            if sub_est is None:
                nones = dict.fromkeys(ids, None)
                new_fits.update(nones)
                if is_transform:
                    if none_passthrough:
                        new_Xs.update(zip(ids, get(ids, Xs)))
                        new_ys.update(zip(ids, get(ids, ys)))
                    else:
                        new_Xs.update(nones)
                        new_ys.update(nones)
            else:
                # Extract the proper subset of Xs, ys
                sub_Xs = get(ids, Xs)
                sub_ys = get(ids, ys)
                # Only subset the parameters/tokens if necessary
                if sub_fields:
                    sub_tokens = list(pluck(sub_inds, get(ids, tokens)))
                    sub_params = list(pluck(sub_inds, get(ids, params)))
                else:
                    sub_tokens = sub_params = None

                if is_transform:
                    sub_fits, sub_Xs, sub_ys = do_fit_transform(
                        dsk,
                        next_token,
                        sub_est,
                        cv,
                        sub_fields,
                        sub_tokens,
                        sub_params,
                        sub_Xs,
                        sub_ys,
                        sub_fit_params,
                        n_splits,
                        error_score,
                    )
                    new_Xs.update(zip(ids, sub_Xs))
                    new_ys.update(zip(ids, sub_ys))
                    new_fits.update(zip(ids, sub_fits))
                else:
                    sub_fits = do_fit(
                        dsk,
                        next_token,
                        sub_est,
                        cv,
                        sub_fields,
                        sub_tokens,
                        sub_params,
                        sub_Xs,
                        sub_ys,
                        sub_fit_params,
                        n_splits,
                        error_score,
                    )
                    new_fits.update(zip(ids, sub_fits))
        # Extract lists of transformed Xs and fit steps
        all_ids = list(range(len(Xs)))
        if is_transform:
            Xs = get(all_ids, new_Xs)
            ys = get(all_ids, new_ys)
        fits = get(all_ids, new_fits)
    elif step is None or isinstance(step, str) and step == "drop":
        # Nothing to do
        fits = [None] * len(Xs)
        if not none_passthrough:
            Xs = fits
    else:
        # Only subset the parameters/tokens if necessary
        if sub_fields:
            sub_tokens = list(pluck(sub_inds, tokens))
            sub_params = list(pluck(sub_inds, params))
        else:
            sub_tokens = sub_params = None

        if is_transform:
            fits, Xs, ys = do_fit_transform(
                dsk,
                next_token,
                step,
                cv,
                sub_fields,
                sub_tokens,
                sub_params,
                Xs,
                ys,
                sub_fit_params,
                n_splits,
                error_score,
            )
        else:
            fits = do_fit(
                dsk,
                next_token,
                step,
                cv,
                sub_fields,
                sub_tokens,
                sub_params,
                Xs,
                ys,
                sub_fit_params,
                n_splits,
                error_score,
            )
    return (fits, Xs, ys) if is_transform else (fits, None, ys)

def _do_pipeline(
    dsk,
    next_token,
    est,
    cv,
    fields,
    tokens,
    params,
    Xs,
    ys,
    fit_params,
    n_splits,
    error_score,
    is_transform,
):
    if "steps" in fields:
        raise NotImplementedError("Setting Pipeline.steps in a gridsearch")

    field_to_index, step_fields_lk = _group_subparams(est.steps, fields)
    fit_params_lk = _group_fit_params(est.steps, fit_params)

    # A list of (step, is_transform)
    instrs = [(s, True) for s in est.steps[:-1]]
    instrs.append((est.steps[-1], is_transform))

    fit_steps = []
    for (step_name, step), transform in instrs:
        fits, Xs, ys = _do_fit_step(
            dsk,
            next_token,
            step,
            cv,
            fields,
            tokens,
            params,
            Xs,
            ys,
            fit_params,
            n_splits,
            error_score,
            step_fields_lk,
            fit_params_lk,
            field_to_index,
            step_name,
            True,
            transform,
        )
        fit_steps.append(fits)

    # Rebuild the pipelines
    step_names = [n for n, _ in est.steps]
    out_ests = []
    out_ests_append = out_ests.append
    name = "pipeline-" + next_token(est)
    m = 0
    seen = {}
    for steps in zip(*fit_steps):
        if steps in seen:
            out_ests_append(seen[steps])
        else:
            for n in range(n_splits):
                dsk[(name, m, n)] = (
                    pipeline,
                    step_names,
                    [None if s is None else s + (n,) for s in steps],
                )
            seen[steps] = (name, m)
            out_ests_append((name, m))
            m += 1

    if is_transform:
        return out_ests, Xs, ys
    return out_ests

def _do_n_samples(dsk, token, Xs, n_splits):
    name = "n_samples-" + token
    n_samples = []
    n_samples_append = n_samples.append
    seen = {}
    m = 0
    for x in Xs:
        if x in seen:
            n_samples_append(seen[x])
        else:
            for n in range(n_splits):
                dsk[name, m, n] = (_num_samples, x + (n,))
            n_samples_append((name, m))
            seen[x] = (name, m)
            m += 1
    return n_samples

def _do_featureunion(
    dsk,
    next_token,
    est,
    cv,
    fields,
    tokens,
    params,
    Xs,
    ys,
    fit_params,
    n_splits,
    error_score,
):
    if "transformer_list" in fields:
        raise NotImplementedError(
            "Setting FeatureUnion.transformer_list " "in a gridsearch"
        )

    (field_to_index, step_fields_lk) = _group_subparams(
        est.transformer_list, fields, ignore=("transformer_weights")
    )
    fit_params_lk = _group_fit_params(est.transformer_list, fit_params)

    token = next_token(est)

    n_samples = _do_n_samples(dsk, token, Xs, n_splits)

    fit_steps = []
    tr_Xs = []
    tr_ys = []
    for (step_name, step) in est.transformer_list:
        fits, out_Xs, out_ys = _do_fit_step(
            dsk,
            next_token,
            step,
            cv,
            fields,
            tokens,
            params,
            Xs,
            ys,
            fit_params,
            n_splits,
            error_score,
            step_fields_lk,
            fit_params_lk,
            field_to_index,
            step_name,
            False,
            True,
        )
        fit_steps.append(fits)
        tr_Xs.append(out_Xs)
        tr_ys.append(out_ys)

    # Rebuild the FeatureUnions
    step_names = [n for n, _ in est.transformer_list]

    if "transformer_weights" in field_to_index:
        index = field_to_index["transformer_weights"]
        weight_lk = {}
        weight_tokens = list(pluck(index, tokens))
        for i, tok in enumerate(weight_tokens):
            if tok not in weight_lk:
                weights = params[i][index]
                if weights is MISSING:
                    weights = est.transformer_weights
                lk = weights or {}
                weight_list = [lk.get(n) for n in step_names]
                weight_lk[tok] = (weights, weight_list)
        weights = get(weight_tokens, weight_lk)
    else:
        lk = est.transformer_weights or {}
        weight_list = [lk.get(n) for n in step_names]
        weight_tokens = repeat(None)
        weights = repeat((est.transformer_weights, weight_list))

    out = []
    out_append = out.append
    fit_name = "feature-union-" + token
    tr_name = "feature-union-concat-" + token
    yt_name = "ytransform-" + token
    m = 0
    seen = {}
    for steps, Xs, ys, wt, (w, wl), nsamp in zip(
        zip(*fit_steps), zip(*tr_Xs), zip(*tr_ys), weight_tokens, weights, n_samples
    ):
        if (steps, wt) in seen:
            out_append(seen[steps, wt])
        else:
            for n in range(n_splits):
                dsk[(fit_name, m, n)] = (
                    feature_union,
                    step_names,
                    ["drop" if s is None else s + (n,) for s in steps],
                    w,
                )
                dsk[(tr_name, m, n)] = (
                    feature_union_concat,
                    [None if x is None else x + (n,) for x in Xs],
                    nsamp + (n,),
                    wl,
                )
                dsk[(yt_name, m, n)] = (
                    feature_union_y,
                    [None if y is None else y + (n,) for y in ys]
                )
            seen[steps, wt] = m
            out_append(m)
            m += 1
    return [(fit_name, i) for i in out], [(tr_name, i) for i in out], [(yt_name, i) for i in out]

# ------------ #
# CV splitting #
# ------------ #

def check_cv(cv=3, y=None, classifier=False):
    """Dask aware version of ``sklearn.model_selection.check_cv``

    Same as the scikit-learn version, but works if ``y`` is a dask object.
    """
    if cv is None:
        cv = 3

    # If ``cv`` is not an integer, the scikit-learn implementation doesn't
    # touch the ``y`` object, so passing on a dask object is fine
    if not is_dask_collection(y) or not isinstance(cv, numbers.Integral):
        return model_selection.check_cv(cv, y, classifier)

    if classifier:
        # ``y`` is a dask object. We need to compute the target type
        target_type = delayed(type_of_target, pure=True)(y).compute()
        if target_type in ("binary", "multiclass"):
            return StratifiedKFold(cv)
    return KFold(cv)

def compute_n_splits(cv, X, y=None, groups=None):
    """Return the number of splits.

    Parameters
    ----------
    cv : BaseCrossValidator
    X, y, groups : array_like, dask object, or None

    Returns
    -------
    n_splits : int
    """
    if not any(is_dask_collection(i) for i in (X, y, groups)):
        return cv.get_n_splits(X, y, groups)

    if isinstance(cv, (_BaseKFold, BaseShuffleSplit)):
        return cv.n_splits

    elif isinstance(cv, PredefinedSplit):
        return len(cv.unique_folds)

    elif isinstance(cv, _CVIterableWrapper):
        return len(cv.cv)

    elif isinstance(cv, (LeaveOneOut, LeavePOut)) and not is_dask_collection(X):
        # Only `X` is referenced for these classes
        return cv.get_n_splits(X, None, None)

    elif isinstance(cv, (LeaveOneGroupOut, LeavePGroupsOut)) and not is_dask_collection(
        groups
    ):
        # Only `groups` is referenced for these classes
        return cv.get_n_splits(None, None, groups)

    else:
        return delayed(cv).get_n_splits(X, y, groups).compute()

def _normalize_n_jobs(n_jobs):
    if not isinstance(n_jobs, int):
        raise TypeError("n_jobs should be an int, got %s" % n_jobs)
    if n_jobs == -1:
        n_jobs = None  # Scheduler default is use all cores
    elif n_jobs < -1:
        n_jobs = cpu_count() + 1 + n_jobs
    return n_jobs

class StaticDaskSearchMixin:
    """Mixin for CV classes that work off a static task graph.

    This is not appropriate for adaptive / incremental training.
    """

    def visualize(self, filename="mydask", format=None, **kwargs):
        """Render the task graph for this parameter search using ``graphviz``.

        Requires ``graphviz`` to be installed.

        Parameters
        ----------
        filename : str or None, optional
            The name (without an extension) of the file to write to disk.  If
            `filename` is None, no file will be written, and we communicate
            with dot using only pipes.
        format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
            Format in which to write output file.  Default is 'png'.
        **kwargs
            Additional keyword arguments to forward to
            ``dask.dot.to_graphviz``.

        Returns
        -------
        result : IPython.diplay.Image, IPython.display.SVG, or None
            See ``dask.dot.dot_graph`` for more information.
        """
        check_is_fitted(self, "dask_graph_")
        return dask.visualize(
            self.dask_graph_, filename=filename, format=format, **kwargs
        )

class DaskBaseSearchCV(BaseEstimator, MetaEstimatorMixin):
    """Base class for hyper parameter search with cross-validation."""

    def __init__(
        self,
        estimator,
        scoring=None,
        iid=True,
        refit=True,
        cv=None,
        error_score="raise",
        return_train_score=_RETURN_TRAIN_SCORE_DEFAULT,
        scheduler=None,
        n_jobs=-1,
        cache_cv=True,
    ):
        self.scoring = scoring
        self.estimator = estimator
        self.iid = iid
        self.refit = refit
        self.cv = cv
        self.error_score = error_score
        self.return_train_score = return_train_score
        self.scheduler = scheduler
        self.n_jobs = n_jobs
        self.cache_cv = cache_cv

    def _check_if_refit(self, attr):
        if not self.refit:
            raise AttributeError(
                "'{}' is not a valid attribute with " "'refit=False'.".format(attr)
            )

    @property
    def _estimator_type(self):
        return self.estimator._estimator_type

    @property
    def best_params_(self):
        check_is_fitted(self, "cv_results_")
        self._check_if_refit("best_params_")
        return self.cv_results_["params"][self.best_index_]

    @property
    def best_score_(self):
        check_is_fitted(self, "cv_results_")
        self._check_if_refit("best_score_")
        if self.multimetric_:
            key = self.refit
        else:
            key = "score"
        return self.cv_results_["mean_test_{}".format(key)][self.best_index_]

    def _check_is_fitted(self, method_name):
        if not self.refit:
            msg = (
                "This {0} instance was initialized with refit=False. {1} "
                "is available only after refitting on the best "
                "parameters."
            ).format(type(self).__name__, method_name)
            raise NotFittedError(msg)
        else:
            check_is_fitted(self, "best_estimator_")

    @property
    def classes_(self):
        self._check_is_fitted("classes_")
        return self.best_estimator_.classes_

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def predict(self, X):
        self._check_is_fitted("predict")
        return self.best_estimator_.predict(X)

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def predict_proba(self, X):
        self._check_is_fitted("predict_proba")
        return self.best_estimator_.predict_proba(X)

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def predict_log_proba(self, X):
        self._check_is_fitted("predict_log_proba")
        return self.best_estimator_.predict_log_proba(X)

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def decision_function(self, X):
        self._check_is_fitted("decision_function")
        return self.best_estimator_.decision_function(X)

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def transform(self, X):
        self._check_is_fitted("transform")
        return self.best_estimator_.transform(X)

    @if_delegate_has_method(delegate=("best_estimator_", "estimator"))
    @derived_from(BaseSearchCV)
    def inverse_transform(self, Xt):
        self._check_is_fitted("inverse_transform")
        return self.best_estimator_.transform(Xt)

    @derived_from(BaseSearchCV)
    def score(self, X, y=None):
        if self.scorer_ is None:
            raise ValueError(
                "No score function explicitly defined, "
                "and the estimator doesn't provide one %s" % self.best_estimator_
            )
        return self.scorer_(self.best_estimator_, X, y)

    def fit(self, X, y=None, groups=None, **fit_params):
        """Run fit with all sets of parameters.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Training vector, where n_samples is the number of samples and
            n_features is the number of features.
        y : array-like, shape = [n_samples] or [n_samples, n_output], optional
            Target relative to X for classification or regression;
            None for unsupervised learning.
        groups : array-like, shape = [n_samples], optional
            Group labels for the samples used while splitting the dataset into
            train/test set.
        **fit_params
            Parameters passed to the ``fit`` method of the estimator
        """
        estimator = self.estimator
        from .._compat import _check_multimetric_scoring

        scorer, multimetric = _check_multimetric_scoring(
            estimator, scoring=self.scoring
        )
        if not multimetric:
            scorer = scorer["score"]
        self.multimetric_ = multimetric

        if self.multimetric_:
            if self.refit is not False and (
                not isinstance(self.refit, str)
                # This will work for both dict / list (tuple)
                or self.refit not in scorer
            ):
                raise ValueError(
                    "For multi-metric scoring, the parameter "
                    "refit must be set to a scorer key "
                    "to refit an estimator with the best "
                    "parameter setting on the whole data and "
                    "make the best_* attributes "
                    "available for that metric. If this is "
                    "not needed, refit should be set to "
                    "False explicitly. %r was ."
                    "passed." % self.refit
                )
        self.scorer_ = scorer

        error_score = self.error_score
        if not (isinstance(error_score, numbers.Number) or error_score == "raise"):
            raise ValueError(
                "error_score must be the string 'raise' or a" " numeric value."
            )

        candidate_params = list(self._get_param_iterator())
        dsk, keys, n_splits, _ = build_cv_graph(
            estimator,
            self.cv,
            self.scorer_,
            candidate_params,
            X,
            y=y,
            groups=groups,
            fit_params=fit_params,
            iid=self.iid,
            error_score=error_score,
            return_train_score=self.return_train_score,
            cache_cv=self.cache_cv,
        )

        n_jobs = _normalize_n_jobs(self.n_jobs)
        scheduler = dask.base.get_scheduler(scheduler=self.scheduler)

        if not scheduler:
            scheduler = dask.threaded.get
        if scheduler is dask.threaded.get and n_jobs == 1:
            scheduler = dask.local.get_sync

        if "Client" in type(getattr(scheduler, "__self__", None)).__name__:
            futures = scheduler(
                dsk, keys, allow_other_workers=True, num_workers=n_jobs, sync=False
            )

            result_map = {}
            ac = as_completed(futures, with_results=True, raise_errors=False)
            for batch in ac.batches():
                for future, result in batch:
                    if future.status == "finished":
                        result_map[future.key] = result
                    else:
                        logger.warning("{} has failed... retrying".format(future.key))
                        future.retry()
                        ac.add(future)

            out = [result_map[k] for k in keys]
        else:
            out = scheduler(dsk, keys, num_workers=n_jobs)

        if self.iid:
            weights = out[0]
            scores = out[1:]
        else:
            weights = None
            scores = out

        if multimetric:
            metrics = list(scorer.keys())
            scorer = self.refit
        else:
            metrics = None
            scorer = "score"

        cv_results = create_cv_results(
            scores, candidate_params, n_splits, error_score, weights, metrics
        )

        results = handle_deprecated_train_score(cv_results, self.return_train_score)
        self.dask_graph_ = dsk
        self.n_splits_ = n_splits
        self.cv_results_ = results

        if self.refit:
            self.best_index_ = np.flatnonzero(
                results["rank_test_{}".format(scorer)] == 1
            )[0]

            best_params = candidate_params[self.best_index_]
            dsk, keys = build_refit_graph(estimator, X, y, best_params, fit_params)

            out = scheduler(dsk, keys, num_workers=n_jobs)
            self.best_estimator_ = out[0]

        return self

_DOC_TEMPLATE = """{oneliner}

{name} implements a "fit" and a "score" method.
It also implements "predict", "predict_proba", "decision_function",
"transform" and "inverse_transform" if they are implemented in the
estimator used.

{description}

Parameters
----------
estimator : estimator object.
    This is assumed to implement the scikit-learn estimator interface.
    Either estimator needs to provide a ``score`` function,
    or ``scoring`` must be passed.

{parameters}

scoring : string, callable, list/tuple, dict or None, default: None
    A single string or a callable to evaluate the predictions on the test
    set.

    For evaluating multiple metrics, either give a list of (unique) strings
    or a dict with names as keys and callables as values.

    NOTE that when using custom scorers, each scorer should return a single
    value. Metric functions returning a list/array of values can be wrapped
    into multiple scorers that return one value each.

    If None, the estimator's default scorer (if available) is used.

iid : boolean, default=True
    If True, the data is assumed to be identically distributed across
    the folds, and the loss minimized is the total loss per sample,
    and not the mean loss across the folds.

cv : int, cross-validation generator or an iterable, optional
    Determines the cross-validation splitting strategy.
    Possible inputs for cv are:
        - None, to use the default 3-fold cross validation,
        - integer, to specify the number of folds in a ``(Stratified)KFold``,
        - An object to be used as a cross-validation generator.
        - An iterable yielding train, test splits.

    For integer/None inputs, if the estimator is a classifier and ``y`` is
    either binary or multiclass, ``StratifiedKFold`` is used. In all
    other cases, ``KFold`` is used.

refit : boolean, or string, default=True
    Refit an estimator using the best found parameters on the whole
    dataset.

    For multiple metric evaluation, this needs to be a string denoting the
    scorer is used to find the best parameters for refitting the estimator
    at the end.

    The refitted estimator is made available at the ``best_estimator_``
    attribute and permits using ``predict`` directly on this
    ``GridSearchCV`` instance.

    Also for multiple metric evaluation, the attributes ``best_index_``,
    ``best_score_`` and ``best_parameters_`` will only be available if
    ``refit`` is set and all of them will be determined w.r.t this specific
    scorer.

    See ``scoring`` parameter to know more about multiple metric
    evaluation.

error_score : 'raise' (default) or numeric
    Value to assign to the score if an error occurs in estimator fitting.
    If set to 'raise', the error is raised. If a numeric value is given,
    FitFailedWarning is raised. This parameter does not affect the refit
    step, which will always raise the error.

return_train_score : boolean, default=True
    If ``'False'``, the ``cv_results_`` attribute will not include training
    scores.

    Note that for scikit-learn >= 0.19.1, the default of ``True`` is
    deprecated, and a warning will be raised when accessing train score results
    without explicitly asking for train scores.

scheduler : string, callable, Client, or None, default=None
    The dask scheduler to use. Default is to use the global scheduler if set,
    and fallback to the threaded scheduler otherwise. To use a different
    scheduler either specify it by name (either "threading", "multiprocessing",
    or "synchronous"), pass in a ``dask.distributed.Client``, or provide a
    scheduler ``get`` function.

n_jobs : int, default=-1
    Number of jobs to run in parallel. Ignored for the synchronous and
    distributed schedulers. If ``n_jobs == -1`` [default] all cpus are used.
    For ``n_jobs < -1``, ``(n_cpus + 1 + n_jobs)`` are used.

cache_cv : bool, default=True
    Whether to extract each train/test subset at most once in each worker
    process, or every time that subset is needed. Caching the splits can
    speedup computation at the cost of increased memory usage per worker
    process.

    If True, worst case memory usage is ``(n_splits + 1) * (X.nbytes +
    y.nbytes)`` per worker. If False, worst case memory usage is
    ``(n_threads_per_worker + 1) * (X.nbytes + y.nbytes)`` per worker.

Examples
--------
{example}

Attributes
----------
cv_results_ : dict of numpy (masked) ndarrays
    A dict with keys as column headers and values as columns, that can be
    imported into a pandas ``DataFrame``.

    For instance the below given table

    +------------+-----------+------------+-----------------+---+---------+
    |param_kernel|param_gamma|param_degree|split0_test_score|...|rank.....|
    +============+===========+============+=================+===+=========+
    |  'poly'    |     --    |      2     |        0.8      |...|    2    |
    +------------+-----------+------------+-----------------+---+---------+
    |  'poly'    |     --    |      3     |        0.7      |...|    4    |
    +------------+-----------+------------+-----------------+---+---------+
    |  'rbf'     |     0.1   |     --     |        0.8      |...|    3    |
    +------------+-----------+------------+-----------------+---+---------+
    |  'rbf'     |     0.2   |     --     |        0.9      |...|    1    |
    +------------+-----------+------------+-----------------+---+---------+

    will be represented by a ``cv_results_`` dict of::

        {{
        'param_kernel': masked_array(data = ['poly', 'poly', 'rbf', 'rbf'],
                                        mask = [False False False False]...)
        'param_gamma': masked_array(data = [-- -- 0.1 0.2],
                                    mask = [ True  True False False]...),
        'param_degree': masked_array(data = [2.0 3.0 -- --],
                                        mask = [False False  True  True]...),
        'split0_test_score'  : [0.8, 0.7, 0.8, 0.9],
        'split1_test_score'  : [0.82, 0.5, 0.7, 0.78],
        'mean_test_score'    : [0.81, 0.60, 0.75, 0.82],
        'std_test_score'     : [0.02, 0.01, 0.03, 0.03],
        'rank_test_score'    : [2, 4, 3, 1],
        'split0_train_score' : [0.8, 0.7, 0.8, 0.9],
        'split1_train_score' : [0.82, 0.7, 0.82, 0.5],
        'mean_train_score'   : [0.81, 0.7, 0.81, 0.7],
        'std_train_score'    : [0.03, 0.04, 0.03, 0.03],
        'mean_fit_time'      : [0.73, 0.63, 0.43, 0.49],
        'std_fit_time'       : [0.01, 0.02, 0.01, 0.01],
        'mean_score_time'    : [0.007, 0.06, 0.04, 0.04],
        'std_score_time'     : [0.001, 0.002, 0.003, 0.005],
        'params'             : [{{'kernel': 'poly', 'degree': 2}}, ...],
        }}

    NOTE that the key ``'params'`` is used to store a list of parameter
    settings dict for all the parameter candidates.

    The ``mean_fit_time``, ``std_fit_time``, ``mean_score_time`` and
    ``std_score_time`` are all in seconds.

best_estimator_ : estimator
    Estimator that was chosen by the search, i.e. estimator
    which gave highest score (or smallest loss if specified)
    on the left out data. Not available if refit=False.

best_score_ : float or dict of floats
    Score of best_estimator on the left out data.
    When using multiple metrics, ``best_score_`` will be a dictionary
    where the keys are the names of the scorers, and the values are
    the mean test score for that scorer.

best_params_ : dict
    Parameter setting that gave the best results on the hold out data.

best_index_ : int or dict of ints
    The index (of the ``cv_results_`` arrays) which corresponds to the best
    candidate parameter setting.

    The dict at ``search.cv_results_['params'][search.best_index_]`` gives
    the parameter setting for the best model, that gives the highest
    mean score (``search.best_score_``).

    When using multiple metrics, ``best_index_`` will be a dictionary
    where the keys are the names of the scorers, and the values are
    the index with the best mean score for that scorer, as described above.

scorer_ : function or dict of functions
    Scorer function used on the held out data to choose the best
    parameters for the model. A dictionary of ``{{scorer_name: scorer}}``
    when multiple metrics are used.

n_splits_ : int
    The number of cross-validation splits (folds/iterations).

Notes
------
The parameters selected are those that maximize the score of the left out
data, unless an explicit score is passed in which case it is used instead.
"""

# ------------ #
# GridSearchCV #
# ------------ #

_grid_oneliner = """\
Exhaustive search over specified parameter values for an estimator.\
"""
_grid_description = """\
The parameters of the estimator used to apply these methods are optimized
by cross-validated grid-search over a parameter grid.\
"""
_grid_parameters = """\
param_grid : dict or list of dictionaries
    Dictionary with parameters names (string) as keys and lists of
    parameter settings to try as values, or a list of such
    dictionaries, in which case the grids spanned by each dictionary
    in the list are explored. This enables searching over any sequence
    of parameter settings.\
"""
_grid_example = """\
>>> import dask_ml.model_selection as dcv
>>> from sklearn import svm, datasets
>>> iris = datasets.load_iris()
>>> parameters = {'kernel': ['linear', 'rbf'], 'C': [1, 10]}
>>> svc = svm.SVC()
>>> clf = dcv.GridSearchCV(svc, parameters)
>>> clf.fit(iris.data, iris.target)  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
GridSearchCV(cache_cv=..., cv=..., error_score=...,
        estimator=SVC(C=..., cache_size=..., class_weight=..., coef0=...,
                      decision_function_shape=..., degree=..., gamma=...,
                      kernel=..., max_iter=-1, probability=False,
                      random_state=..., shrinking=..., tol=...,
                      verbose=...),
        iid=..., n_jobs=..., param_grid=..., refit=..., return_train_score=...,
        scheduler=..., scoring=...)
>>> sorted(clf.cv_results_.keys())  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
['mean_fit_time', 'mean_score_time', 'mean_test_score',...
 'mean_train_score', 'param_C', 'param_kernel', 'params',...
 'rank_test_score', 'split0_test_score',...
 'split0_train_score', 'split1_test_score', 'split1_train_score',...
 'split2_test_score', 'split2_train_score',...
 'std_fit_time', 'std_score_time', 'std_test_score', 'std_train_score'...]\
"""

class GridSearchCV(StaticDaskSearchMixin, DaskBaseSearchCV):
    __doc__ = _DOC_TEMPLATE.format(
        name="GridSearchCV",
        oneliner=_grid_oneliner,
        description=_grid_description,
        parameters=_grid_parameters,
        example=_grid_example,
    )

    def __init__(
        self,
        estimator,
        param_grid,
        scoring=None,
        iid=True,
        refit=True,
        cv=None,
        error_score="raise",
        return_train_score=_RETURN_TRAIN_SCORE_DEFAULT,
        scheduler=None,
        n_jobs=-1,
        cache_cv=True,
    ):
        super(GridSearchCV, self).__init__(
            estimator=estimator,
            scoring=scoring,
            iid=iid,
            refit=refit,
            cv=cv,
            error_score=error_score,
            return_train_score=return_train_score,
            scheduler=scheduler,
            n_jobs=n_jobs,
            cache_cv=cache_cv,
        )

        _check_param_grid(param_grid)
        self.param_grid = param_grid

    def _get_param_iterator(self):
        """Return ParameterGrid instance for the given param_grid"""
        return model_selection.ParameterGrid(self.param_grid)

# ------------------ #
# RandomizedSearchCV #
# ------------------ #

_randomized_oneliner = "Randomized search on hyper parameters."
_randomized_description = """\
In contrast to GridSearchCV, not all parameter values are tried out, but
rather a fixed number of parameter settings is sampled from the specified
distributions. The number of parameter settings that are tried is
given by n_iter.

If all parameters are presented as a list, sampling without replacement is
performed. If at least one parameter is given as a distribution, sampling
with replacement is used. It is highly recommended to use continuous
distributions for continuous parameters.\
"""
_randomized_parameters = """\
param_distributions : dict
    Dictionary with parameters names (string) as keys and distributions
    or lists of parameters to try. Distributions must provide a ``rvs``
    method for sampling (such as those from scipy.stats.distributions).
    If a list is given, it is sampled uniformly.

n_iter : int, default=10
    Number of parameter settings that are sampled. n_iter trades
    off runtime vs quality of the solution.

random_state : int or RandomState
    Pseudo random number generator state used for random uniform sampling
    from lists of possible values instead of scipy.stats distributions.\
"""
_randomized_example = """\
>>> import dask_ml.model_selection as dcv
>>> from scipy.stats import expon
>>> from sklearn import svm, datasets
>>> iris = datasets.load_iris()
>>> parameters = {'C': expon(scale=100), 'kernel': ['linear', 'rbf']}
>>> svc = svm.SVC()
>>> clf = dcv.RandomizedSearchCV(svc, parameters, n_iter=100)
>>> clf.fit(iris.data, iris.target)  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
RandomizedSearchCV(cache_cv=..., cv=..., error_score=...,
        estimator=SVC(C=..., cache_size=..., class_weight=..., coef0=...,
                      decision_function_shape=..., degree=..., gamma=...,
                      kernel=..., max_iter=..., probability=...,
                      random_state=..., shrinking=..., tol=...,
                      verbose=...),
        iid=..., n_iter=..., n_jobs=..., param_distributions=...,
        random_state=..., refit=..., return_train_score=...,
        scheduler=..., scoring=...)
>>> sorted(clf.cv_results_.keys())  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
['mean_fit_time', 'mean_score_time', 'mean_test_score',...
 'mean_train_score', 'param_C', 'param_kernel', 'params',...
 'rank_test_score', 'split0_test_score',...
 'split0_train_score', 'split1_test_score', 'split1_train_score',...
 'split2_test_score', 'split2_train_score',...
 'std_fit_time', 'std_score_time', 'std_test_score', 'std_train_score'...]\
"""

class RandomizedSearchCV(StaticDaskSearchMixin, DaskBaseSearchCV):
    __doc__ = _DOC_TEMPLATE.format(
        name="RandomizedSearchCV",
        oneliner=_randomized_oneliner,
        description=_randomized_description,
        parameters=_randomized_parameters,
        example=_randomized_example,
    )

    def __init__(
        self,
        estimator,
        param_distributions,
        n_iter=10,
        random_state=None,
        scoring=None,
        iid=True,
        refit=True,
        cv=None,
        error_score="raise",
        return_train_score=_RETURN_TRAIN_SCORE_DEFAULT,
        scheduler=None,
        n_jobs=-1,
        cache_cv=True,
    ):
        super(RandomizedSearchCV, self).__init__(
            estimator=estimator,
            scoring=scoring,
            iid=iid,
            refit=refit,
            cv=cv,
            error_score=error_score,
            return_train_score=return_train_score,
            scheduler=scheduler,
            n_jobs=n_jobs,
            cache_cv=cache_cv,
        )

        self.param_distributions = param_distributions
        self.n_iter = n_iter
        self.random_state = random_state

    def _get_param_iterator(self):
        """Return ParameterSampler instance for the given distributions"""
        return model_selection.ParameterSampler(
            self.param_distributions, self.n_iter, random_state=self.random_state
        )


import numpy as np
from functools import lru_cache
import random
from senseTk.common import TrackSet, Det
import torch.optim.lr_scheduler as lr_scd
import torch.optim as optim
from torch.nn.parallel.data_parallel import DataParallel
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.clip_grad import clip_grad_norm_
import torch.nn.functional as F
from torch import nn
import torch
import os
from typing import *
from collections import deque
from src.lib.datasets.DetApprDataset import PseudoTracks2, PseudoTracks3, RealTracks, SparseRTracks as SRealTracks
from src.lib.models.networks.THN import THN
from src.lib.models.networks.THN2 import THN as THN2

# from src.lib.models.networks.TrackMem import TrackMem
# from src.lib.models.networks.StructDNC import StructDNC as TrackMem

def to_device(data: Any, device: Union[torch.device, AnyStr] = 'cuda') -> Any:
    if isinstance(data, torch.Tensor):
        # print(data.shape)
        return data.to(device)
    if isinstance(data, Mapping):
        # print('map:', [k for k in data])
        return {k: to_device(v, device=device) for k, v in data.items()}
    if isinstance(data, list):
        # print('list:')
        return [to_device(i, device=device) for i in data]
    if isinstance(data, tuple):
        # print('list:')
        return tuple(to_device(i, device=device) for i in data)
    return data

history = {}

def parse_losses(out):
    global history
    if not isinstance(out, list):
        out = [out]
    keys = set()
    for one in out:
        keys |= one.keys()
    cnts = {k: 0 for k in keys}
    items = {k: 0 for k in keys}
    loss = 0
    log = {}
    for one in out:
        for k in one:
            if k.startswith('loss') and one[k] < 1e-8:
                continue
            items[k] += one[k]
            cnts[k] += 1
    for k in sorted(keys):
        if k.startswith('loss'):
            single_loss = items[k] / max(cnts[k], 1)
            loss += single_loss
            log[k] = float(single_loss)
        else:
            log[k] = float(items[k]) / max(cnts[k], 1)
        if k not in history:
            history[k] = deque(maxlen=1000)
        if k.startswith('loss') and log[k] < 1e-8:
            continue
        history[k].append(log[k])
    logh = {k: sum(history[k]) / max(len(history[k]), 1) for k in log}
    return loss, log, logh

def get_frame_data(data, frame_ind):
    meta = ['origin_size']
    frame_data = {
        k: ((v[:, frame_ind] if isinstance(v, torch.Tensor)
             else v[frame_ind]) if k not in meta else v)
        for k, v in data.items()
    }
    # frame_data['dets'] = frame_data['dets'][frame_data['inds']]
    # frame_data['feats'] = frame_data['feats'][frame_data['inds']]
    # frame_data['confs'] = frame_data['confs'][frame_data['inds']]
    # frame_data['gt_dets'] = frame_data['gt_dets'][frame_data['gt_inds']]
    # frame_data['gt_ids'] = frame_data['gt_ids'][frame_data['gt_inds']]
    # frame_data['gt_vis'] = frame_data['gt_vis'][frame_data['gt_inds']]
    return frame_data

def load_model(path, model):
    a = torch.load(path, map_location='cpu')
    start_epoch = a.get('start_epoch', 0)
    opt_state = a.get('opt_state', None)
    if 'state_dict' in a:
        model.load_state_dict(a['state_dict'])
    else:
        model.load_state_dict(a)
    return start_epoch, model, opt_state

p4_root = '/home/toka/data/P4data'
gt_root = '/mnt/lustre/share/fengweitao/MOT17/images/train'
seqs = ['MOT17-%02d-SDP' % i for i in [2, 4, 5, 9, 10, 11, 13]]

if __name__ == '__main__':
    epoch_outer = 30
    epoch_inner = 30
    mile_stone = 1
    track_step = 2
    batch_size = 8
    init_lr = 0.01
    devices = ['cuda:0']  # , 'cuda:1']  # , 'cuda:2', 'cuda:3']
    model = THN2(state_len=16)
    print(model)
    model.init_weights()
    start_epoch = 0
    opt_state = None
    # start_epoch, model, opt_state = load_model(
    #     '/home/toka/code/FairMOT/chkpts/thn2/epoch_7_26.pth', model)
    device = 'cuda:0'
    model.to_device(device)
    model = model.to(device)
    model = DataParallel(model, devices)
    ps = SRealTracks(
        [os.path.join(p4_root, '%s_dense.p4data' % seq)
            for seq in seqs], '/home/toka/data/P4data/seqinfos.yaml',
        [os.path.join(gt_root, seq, 'gt', 'gt.txt') for seq in seqs], clip_len=epoch_outer * track_step)
    dl = DataLoader(ps, batch_size=batch_size, shuffle=True,
                    num_workers=8, drop_last=True)
    base_parameters = (y for x, y in model.named_parameters())
    opt = optim.SGD(base_parameters, lr=init_lr, weight_decay=1e-4)
    if opt_state is not None:
        opt.load_state_dict(opt_state)
    lr_scheduler = lr_scd.MultiStepLR(opt, [50, 600], gamma=0.1)
    if start_epoch > 0:
        lr_scheduler.step(epoch=start_epoch)
    start_outter = start_epoch // epoch_inner
    start_inner = start_epoch % epoch_inner
    for eo in range(start_outter, epoch_outer):
        # second, third = max(eo - 1, 0) * track_step, (eo + 1) * track_step
        second, third = (eo + 1 - mile_stone) * \
            track_step, (eo + 2 - mile_stone) * track_step
        for e in range(start_inner, epoch_inner):
            for i, data in enumerate(dl):
                if eo < mile_stone:
                    second = 0
                    third = 2
                else:
                    pre_run = np.random.randint(
                        (eo + 1 - mile_stone) * track_step + 1)
                    second = int(pre_run)
                    third = second + track_step
                print_flag = i < epoch_outer or i % 10 == 1
                trk_state = model.module.init_state(batch_size=batch_size)
                data = to_device(data, device=device)
                model.eval()
                for j in range(second):
                    frame_data = get_frame_data(data, j)
                    with torch.no_grad():
                        out, trk_state = model(frame_data, state=trk_state)
                        trk_state = model.module.map_gt_ids(
                            out, frame_data, state=trk_state)
                model.train()
                opt.zero_grad()
                losses = []
                for j in range(second, third):
                    frame_data = get_frame_data(data, j)
                    # print(model._cs_pos[0].weight.mean(), 'before')
                    output, trk_state = model(frame_data, state=trk_state)
                    # input()
                    trk_state = model.module.map_gt_ids(
                        output, frame_data, state=trk_state)
                    loss_j = model.module.loss(
                        output, frame_data, verbose=print_flag)
                    losses.append(loss_j)
                    # print(model._cs_pos[0].weight.mean(), 'after')
                loss, logs, logh = parse_losses(losses)
                if isinstance(loss, torch.Tensor):
                    # for name, one in model.named_parameters():
                    #     if one.grad is not None:
                    #         print(name, one.grad.mean())
                    loss.backward()
                    # for name, one in model.named_parameters():
                    #     if one.grad is not None:
                    #         print(name, one.grad.mean())
                    clip_grad_norm_(model.parameters(), 1)
                    opt.step()
                # print(model._cs_pos[0].weight.mean(), 'stepped')
                lr = lr_scheduler.get_last_lr()[0]
                if print_flag:
                    print('epoch %d (outer %d inner %d [%d:%d]) iter %d/%d: lr %.6f' %
                          (lr_scheduler.last_epoch, eo, e, second, third, i, len(dl),
                           float(lr)))
                    print('\t', end='')
                    for k in logs:
                        print('.%s: %.6f(%.6f), ' % (k, logs[k], logh[k]),
                              end='')
                    print(' loss %.6f' % float(loss))
                    for j in range(len(losses)):
                        print('\tstep %d;' % j, end='')
                        for k in losses[j]:
                            if k.startswith('loss.'):
                                print('.%s %.6f ' % (k, float(losses[j][k])),
                                      end='')
                        print()
                        if j == len(losses) - 1:
                            topks, topki = output['thn_scores'][0].flatten().topk(
                                5)
                            print(topks)
                            print(torch.gather(
                                output['pos_sc'][0].flatten(), 0, topki))
                            print(torch.gather(
                                output['apr_sc'][0].flatten(), 0, topki))
                            print(torch.gather(
                                output['det_sc'][0].flatten(), 0, topki))
                            print(torch.gather(
                                output['pre_sc'][0].flatten(), 0, topki))
                            print(model.module._gt_ids_for_pred[0][topki])
                            print(output['thn_act_scores'][0][topki])
                            print(model.module._iter_cnt, model.module._device)
                    # for k in output:
                    #     if isinstance(output[k], torch.Tensor):
                    #         print(k, output[k].shape)
                    for hypn in ['_pos_alpha', '_pos_beta', '_apr_alpha', '_apr_beta', '_conf_lambda', '_act_conf_lambda']:
                        hyp = getattr(model.module, hypn)
                        print('[%s]%.6f' % (hypn, float(hyp.data)), end=' ')
                    print()
            lr_scheduler.step()
            if (e + 1) % epoch_inner == 0 or (e + 1) % 2 == 0:
                torch.save({'state_dict': model.module.state_dict(), 'start_epoch': lr_scheduler.last_epoch, 'optim': opt.state_dict()},
                           'chkpts/thn2/epoch_%d_%d.pth' % (eo + 1, e + 1))


import numpy as np
import cv2
import os
import gc

from .face_utils import get_tracks, extract_sequence, extract_face
from .models import *

PATH_PREFIX = '/weights/ntech/'
DETECTOR_WEIGHTS_PATH = os.path.join(PATH_PREFIX, 'WIDERFace_DSFD_RES152.fp16.pth')
VIDEO_SEQUENCE_MODEL_WEIGHTS_PATH = os.path.join(PATH_PREFIX, 'efficientnet-b7_ns_seq_aa-original-mstd0.5_100k_v4_cad79a/snapshot_100000.fp16.pth')
FIRST_VIDEO_FACE_MODEL_WEIGHTS_PATH = os.path.join(PATH_PREFIX, 'efficientnet-b7_ns_aa-original-mstd0.5_large_crop_100k_v4_cad79a/snapshot_100000.fp16.pth')
SECOND_VIDEO_FACE_MODEL_WEIGHTS_PATH = os.path.join(PATH_PREFIX, 'efficientnet-b7_ns_aa-original-mstd0.5_re_100k_v4_cad79a/snapshot_100000.fp16.pth')

VIDEO_FACE_MODEL_TRACK_STEP = 2
VIDEO_SEQUENCE_MODEL_SEQUENCE_LENGTH = 7
VIDEO_SEQUENCE_MODEL_TRACK_STEP = 14

class video_reader:
    def __init__(self, face_detector, video_target_fps=15, detector_step=3):
        self.video_target_fps = video_target_fps
        self.face_detector = face_detector
        self.detector_step = detector_step

    def process_video(self, video_path):
        sample = {
            'frames': [],
            'tracks': []
        }
        capture = cv2.VideoCapture(video_path)
        frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if frame_count == 0:
            return sample

        fps = int(capture.get(cv2.CAP_PROP_FPS))
        video_step = round(fps / self.video_target_fps)
        if video_step == 0:
            return sample

        for i in range(frame_count):
            capture.grab()
            if i % video_step != 0:
                continue
            ret, frame = capture.retrieve()
            if not ret:
                continue

            sample['frames'].append(frame)

        detector_frames = sample['frames'][::self.detector_step]
        DETECTOR_BATCH_SIZE = 16
        detections = []
        for start in range(0, len(detector_frames), DETECTOR_BATCH_SIZE):
            end = min(len(detector_frames), start + DETECTOR_BATCH_SIZE)
            detections_batch = self.face_detector.detect(
                detector_frames[start:end])
            for detections_per_frame in detections_batch:
                detections.append({key: value.cpu().numpy()
                                   for key, value in detections_per_frame.items()})
        sample['tracks'] = get_tracks(detections)
        return sample

class Ensemble:
    def __init__(self):
        detector = Detector(DETECTOR_WEIGHTS_PATH)
        self.track_sequences_classifier = TrackSequencesClassifier(
            VIDEO_SEQUENCE_MODEL_WEIGHTS_PATH)
        self.track_faces_classifier = TrackFacesClassifier(
            FIRST_VIDEO_FACE_MODEL_WEIGHTS_PATH, SECOND_VIDEO_FACE_MODEL_WEIGHTS_PATH)
        self.reader = video_reader(detector)

    def inference(self, video_path):
        sample = self.reader.process_video(video_path)
        frames, tracks = sample['frames'], sample['tracks']
        if len(frames) == 0 or len(tracks) == 0:
            return 0.5
        sequence_track_scores = []
        for track in tracks:
            track_sequences = []
            for i, (start_idx, _) in enumerate(
                    track[:-VIDEO_SEQUENCE_MODEL_SEQUENCE_LENGTH + 1:VIDEO_SEQUENCE_MODEL_TRACK_STEP]):
                assert start_idx >= 0 and start_idx + \
                    VIDEO_SEQUENCE_MODEL_SEQUENCE_LENGTH <= len(frames)
                _, bbox = track[i * VIDEO_SEQUENCE_MODEL_TRACK_STEP +
                                VIDEO_SEQUENCE_MODEL_SEQUENCE_LENGTH // 2]
                track_sequences.append(extract_sequence(
                    frames, start_idx, bbox, i % 2 == 0))
            sequence_track_scores.append(
                self.track_sequences_classifier.classify(track_sequences))
        face_track_scores = []
        for track in tracks:
            track_faces = []
            for i, (frame_idx, bbox) in enumerate(track[::VIDEO_FACE_MODEL_TRACK_STEP]):
                face = extract_face(frames[frame_idx], bbox, i % 2 == 0)
                track_faces.append(face)
            face_track_scores.append(
                self.track_faces_classifier.classify(track_faces))

        sequence_track_scores = np.concatenate(sequence_track_scores)
        face_track_scores = np.concatenate(face_track_scores)
        track_probs = np.concatenate(
            (sequence_track_scores, face_track_scores))

        delta = track_probs - 0.5
        sign = np.sign(delta)
        pos_delta = delta > 0
        neg_delta = delta < 0
        track_probs[pos_delta] = np.clip(
            0.5 + sign[pos_delta] * np.power(abs(delta[pos_delta]), 0.65), 0.01, 0.99)
        track_probs[neg_delta] = np.clip(
            0.5 + sign[neg_delta] * np.power(abs(delta[neg_delta]), 0.65), 0.01, 0.99)
        weights = np.power(abs(delta), 1.0) + 1e-4
        video_score = float((track_probs * weights).sum() / weights.sum())

        return video_score

    def __del__(self):
        del self.track_sequences_classifier
        del self.track_faces_classifier
        del self.reader

        gc.collect()


# -*- coding: UTF-8 -*-
# @Time    : 04/08/2020 18:55
# @Author  : BubblyYi、QYD
# @FileName: centerline_net.py
# @Software: PyCharm
import numpy as np
from setting import src_array, spacing, seeds_model, ostia_model, device
from utils import data_preprocess, resample, crop_heart
import torch

def search_seeds_ostias(max_size=(200, 10)):
    '''
    find seeds points arr and ostia points arr
    :param max_size: The first max_size[0] seed points and the first max_size[1] ostia points were selected
    :return: 种子点list和入口点list
    '''
    print("search seeds and ostias")
    spacing_x = spacing[0]
    spacing_y = spacing[1]
    spacing_z = spacing[2]

    re_spacing_img, curr_spacing, resize_factor = resample(src_array, np.array([spacing_z, spacing_x, spacing_y]),
                                                           np.array([1, 1, 1]))
    re_spacing_img, meam_minc, mean_minr, mean_maxc, mean_maxr = crop_heart(re_spacing_img)
    cut_size = 9
    res_seeds = {}
    res_ostia = {}
    count = 0
    random_point_size = 80000
    batch_size = 1000
    new_patch_list = []
    center_coord_list = []
    z, h, w = re_spacing_img.shape
    offset_size = 10
    x_list = np.random.random_integers(meam_minc - offset_size, mean_maxc + offset_size, (random_point_size, 1))
    y_list = np.random.random_integers(mean_minr - offset_size, mean_maxr + offset_size, (random_point_size, 1))
    z_list = np.random.random_integers(0, z, (random_point_size, 1))

    index = np.concatenate([x_list, y_list, z_list], axis=1)
    index = list(set(tuple(x) for x in index))
    for i in index:
        center_x_pixel = i[0]
        center_y_pixel = i[1]
        center_z_pixel = i[2]
        left_x = center_x_pixel - cut_size
        right_x = center_x_pixel + cut_size
        left_y = center_y_pixel - cut_size
        right_y = center_y_pixel + cut_size
        left_z = center_z_pixel - cut_size
        right_z = center_z_pixel + cut_size
        if left_x >= 0 and right_x < h and left_y >= 0 and right_y < w and left_z >= 0 and right_z < z:
            new_patch = np.zeros((cut_size * 2 + 1, cut_size * 2 + 1, cut_size * 2 + 1))
            for ind in range(left_z, right_z + 1):
                src_temp = re_spacing_img[ind].copy()
                new_patch[ind - left_z] = src_temp[left_y:right_y + 1, left_x:right_x + 1]
            count += 1
            input_data = data_preprocess(new_patch)
            new_patch_list.append(input_data)
            center_coord_list.append((center_x_pixel, center_y_pixel, center_z_pixel))
            if count % batch_size == 0:
                input_data = torch.cat(new_patch_list, axis=0)
                inputs = input_data.to(device)
                seeds_outputs = seeds_model(inputs.float())
                seeds_outputs = seeds_outputs.view((len(input_data)))  # view
                seeds_proximity = seeds_outputs.cpu().detach().numpy()
                ostia_outputs = ostia_model(inputs.float())
                ostia_outputs = ostia_outputs.view(len(input_data))
                ostia_proximity = ostia_outputs.cpu().detach().numpy()
                for i in range(batch_size):
                    res_seeds[center_coord_list[i]] = seeds_proximity[i]
                    res_ostia[center_coord_list[i]] = ostia_proximity[i]
                new_patch_list.clear()
                center_coord_list.clear()
                del input_data
                del inputs
                del seeds_outputs
                del ostia_outputs

    positive_count = 0
    for i in res_seeds.values():
        if i > 0:
            positive_count += 1
    res_seeds = sorted(res_seeds.items(), key=lambda item: item[1], reverse=True)
    res_ostia = sorted(res_ostia.items(), key=lambda item: item[1], reverse=True)
    res_seeds = res_seeds[:max_size[0]]
    res_ostia = res_ostia[:max_size[1]]
    return res_seeds, res_ostia


#!/usr/bin/python
# encoding: utf-8
M=float
t=str
X=list
g=False
b=dict
mV=True
mE=print
mF=any
import os
G=os.path
import sys
u=sys.path
u.append('/Users/luoyonggui/PycharmProjects/mayiutils_n1/mayiutils/db')
from pymongo_wrapper import PyMongoWrapper
import pandas as pd
N=pd.date_range
R=pd.to_datetime
L=pd.read_pickle
r=pd.concat
w=pd.DataFrame
U=pd.read_excel
import numpy as np
l=np.where
class Publisher:
 def __init__(x):
  pass
 def register(x):
  pass
 def unregister(x):
  pass
 def notify_all(x):
  pass
class Subscriber:
 def __init__(x):
  pass
 def notify(x):
  pass
from attr import attrs,attrib,fields
from cattr import unstructure,structure
from datetime import datetime,date
c=datetime.strptime
e=datetime.now
@attrs
class Transaction:
 m=attrib(type=M,default=0)
 V=attrib(type=M,default=0)
 E=attrib(type=M,default=0)
 F=attrib(type=datetime,default=e())
 T=attrib(type=t,default='cash')
 k=attrib(type=t,default='pingan')
@attrs
class Trade(Publisher):
 p=attrib(factory=X,repr=g)
 a=attrib(factory=X)
 def register(x,k):
  if k not in x._accounts:
   x._accounts.append(k)
 def unregister(x,k):
  if k not in x._accounts:
   x._accounts.remove(k)
 def notify_all(x):
  for a in x._accounts:
   a.notify(x._transaction_list)
@attrs
class Account(Subscriber):
 d=attrib(type=t)
 S=attrib(factory=X,repr=g)
 W=attrib(factory=b)
 q=attrib(type=M,default=0)
 def notify(x,f):
  for t in f:
   if t.account==x._name:
    x._trade_list.append(t)
    if t.target=='cash':
     x._cash+=t.num
    else:
     m=t.num
     o=t.num*t.price
     h=5 if o*0.00025<5 else o*0.00025
     if t.target in x._positions:
      m+=x._positions[t.target][0]
      o+=x._positions[t.target][1]*x._positions[t.target][0]
     z=(o+h)/m
     x._positions[t.target]=(m,z)
def add_transactions():
 df=U('data/transactions.xlsx')
 r=[]
 for I in df.itertuples():
  if I.target!='cash':
   r.append([I.trade_date,I.account,'cash',1,-1*(I.price*I.num*1+I.fee_rate),0])
 A=w(r,columns=df.columns)
 df=r([df,A],ignore_index=mV)
 mE(df)
 if not df.empty:
  K=PyMongoWrapper()
  K.insertDataframe(df,'finance','transactions')
 else:
  mE('no transactions!')
H=G.dirname(G.abspath(__file__))
j=L(G.join(H,'data/stock_dict.pkl'))
y=L(G.join(H,'data/fund_otc_series.pkl'))
i=j.append(y)
def get_last_day_market_val(T,D):
 if T=='cash':
  return e().date(),1
 if mF([T.endswith('.SZ'),T.endswith('.SH')]):
  s=T
 else:
  s=i.loc[T]
 K=PyMongoWrapper()
 v=K.getCollection('finance',s)
 r=X(K.findAll(v,{'trade_date':{'$lte':c(D,'%Y%m%d')}},fieldlist=['trade_date','close','pct_chg'],sort=[('trade_date',-1)],limit=1))[0]
 return r['trade_date'].date(),r['close']
def get_account_val(J,D):
 J=J[:D]
 Q=J.groupby(['account','target'])['num','cost'].sum()
 Q.reset_index(inplace=mV)
 Q['price']=0
 Q['tdate']=0
 for i in J['target'].unique():
  Q.loc[Q.target==i,['tdate','price']]=get_last_day_market_val(i,D)
 Q['type']=l((Q.account=='PA')&(Q.target!='cash'),'STOCK',l((Q.account=='ZLT')&(Q.target!='cash'),'FUND',l((Q.account.isin(['TTJ','TTA']))&(Q.target!='cash'),'OTC_FUND','cash')))
 Q['market_val']=Q.num*Q.price
 Q['profit']=Q['market_val']-Q['cost']
 s=Q.groupby('account')['market_val'].sum()
 s.loc['TOTAL']=s.sum()
 B=w(s).T
 B.index.name='date'
 B.index=R([D])
 return B
if __name__=='__main__':
 f=[]
 K=PyMongoWrapper()
 v=K.getCollection('finance','transactions')
 n='trade_date account target price num fee_rate'.split()
 rs=K.findAll(v,fieldlist=n,sort=[('trade_date',1)])
 J=w(rs)
 J.set_index('trade_date',inplace=mV)
 J['cost']=J['price']*J['num']*(1+J.fee_rate)
 C=w()
 P=N('20190921','20190926',freq='B')
 for d in P:
  d=d.strftime('%Y%m%d')
  if C.empty:
   C=get_account_val(J,d)
  else:
   C=r([C,get_account_val(J,d)])
 mE(C)
 C.plot(subplots=mV,figsize=(8,14))
# Created by pyminifier (https://github.com/liftoff/pyminifier)