from mhr.config import *


class MHRVedioCuter:
    def __init__(self, speed_ratio=1):
        ## config
        #self.part_pos = {
        #    'pos':[(628,350),(993,565)],
        #    'page': [(781,575),(848,600)],
        #    'hole': [(1167,197),(1250,227)],
        #    'skill': [(1010,260),(1254,600)],
        #}
        self.part_pos = {
            'hole': [1166,200,28,26],
            'skill': [1014,264,240,50],
        }
        self.speed_ratio = speed_ratio


    def iter(self, name):
        vc = cv.VideoCapture(name)
        fps = vc.get(cv.CAP_PROP_FPS)
        print("vedio:", vc.isOpened(), fps)
        label = "00:00:{:05.2f}({})"
        rval = True
        idx=0
        while rval: 
            rval, img = vc.read()
            idx+=1
            if rval and idx%self.speed_ratio == 0:
                yield self._cut_whole(img), self._cut_hole(img), self._cut_skill(img), label.format(idx/fps, idx)

        vc.release()

    def _cut_whole(self, img):
        #pos_w, pos_h, w, h = self.part_pos['skill']
        return img

    def _cut_hole(self, img):
        pos_w, pos_h, w, h = self.part_pos['hole']
        return [ img[pos_h:pos_h+h, pos_w+w*i:pos_w+w*i+w] for i in range(3) ]

    def _cut_skill(self, img):
        pos_w, pos_h, w, h = self.part_pos['skill']
        return [ img[pos_h+h*i:pos_h+h*i+h, pos_w:pos_w+w] for i in range(7) ]



class MHRStoneRecognizeMgr:
    def __init__(self, whole_pkl, hole_pkl, vedio_cutter):
        self.mapping_hole = [0,2,1,3,4]
        cp = torch.load(whole_pkl)
        self.whole_model = GaborFeatureNet(num_classes=2)
        self.whole_model.load_state_dict(cp['model'])
        if torch.cuda.is_available():
            self.whole_model = self.whole_model.cuda()
        #self.whole_model = torch.load(whole_pkl)
        self.hole_feat_model = GaborFeatureGen(0)
        if torch.cuda.is_available():
            self.hole_feat_model = self.hole_feat_model.cuda()
        self.hole_model = pickle.load(open(hole_pkl, 'rb'))
        self.skill_model = MyTrRecognizeNet(image_padding=2)
        self._vedio_cutter = vedio_cutter

    def recognize_image(self, data):
        print(type(data))
        data_whole = tsfm_whole4cv(data[0])
        data_whole = data_whole.unsqueeze(0)
        if torch.cuda.is_available():
            data_whole = data_whole.cuda()
        ret = self.whole_model(data_whole)
        if ret[0][1] - ret[0][0] < 2:
            return False, []
        #new hole
        data_hole = torch.cat([ tsfm_hole4cv(item).unsqueeze(0) for item in data[1] ], dim=0)
        if torch.cuda.is_available():
            data_hole = data_hole.cuda()
        output = self.hole_feat_model(data_hole)
        df = pd.DataFrame(output.tolist())
        ret_hole = list(self.hole_model.predict(df))
        #new skill
        data_skill = torch.cat([ tsfm_skill4cv(item).unsqueeze(0) for item in data[2] ], dim=0)
        ret = self.skill_model(data_skill)
        ret_skill = [ (x[0][0], x[1][0][-1]) for x in filter(lambda sk: sk[0][1] > 0.9 and sk[1][1] > 0.9, ret) ]
        # reuslt
        result = [data[3], ret_hole, ret_skill]
        return True, result

        if len(results) > 0 and dump(results[-1]) == dump(result):
            return 

    def recognize(self, vname, fname=None):
        def dump(rr):
            return '_'.join([ str(x) for x in rr[1] ]) + "|" + '|'.join([ x[0]+":"+x[1] for x in rr[2] ])
        results = []
        i=0
        for data in self._vedio_cutter.iter(vname):
            ok, result = self.recognize_image(data)
            if not ok or (len(results) > 0 and dump(results[-1]) == dump(result)):
                continue
            results.append(result)
        if fname:
            with open(fname, 'w') as f:
                for result in results:
                    line = result[0]
                    line += ','
                    line += ','.join([ str(x) for x in result[1] ])
                    line += ','
                    line += ','.join([ x[0]+","+x[1] for x in result[2] ])
                    line += '\n'
                    f.write(line)
        return results



#mgr = MHRStoneRecognizeMgr(
#    whole_pkl = whole_pkl_file,
#    hole_pkl = hole_pkl_file,
#    vedio_cutter = MHRVedioCuter(speed_ratio),
#)