import json
import math
import torch
import random
from flowvel.calculator import getFootPoint, getdis2, getdis3, cross, getmul, getsec, getrealxy, getH, getrealxybyH
import os.path as osp
import yaml


class camParam:
    def __init__(self, matrix, max, min):
        self.matrix = matrix
        self.max = max
        self.min = min
        self.fReprojErr = 0

    def getInsRandom(self):
        matrix = torch.zeros(1, 9)
        for i in range(9):
            matrix[0][i] = self.min[0][i] + (self.max[0][i]-self.min[0][i])*random.random()
        return camParam(matrix, self.max, self.min)

    def setReprojErr(self, fReprojErr):
        self.fReprojErr = fReprojErr

    def getReprojErr(self):
        return self.fReprojErr

    def getH(self):
        return self.matrix


def get_parameters_auto(v1, v2, line, width, w, h):
    oi = [w / 2, h / 2]  # 画幅中心
    ci = [0, 0]
    v1 = [v1[0] - oi[0], v1[1] - oi[1]]
    v2 = [v2[0] - oi[0], v2[1] - oi[1]]

    vi = getFootPoint(ci, v1, v2)  # 求vi

    ocvi = math.sqrt(getdis2(v1, vi) * getdis2(vi, v2))

    f = math.sqrt(math.fabs((ocvi ** 2) - (getdis2(ci, vi) ** 2)))  # 原码
    K = [[f, 0, oi[0]], [0, f, oi[1]], [0, 0, 1]]  # 得到矩阵K

    len_ocv1 = getdis3([v1[0], v1[1], f], [0, 0, 0])  # Xc = ||ocvi||
    len_ocv2 = getdis3([v2[0], v2[1], f], [0, 0, 0])  # Yc = ||ocv2||

    x = [v1[0] / len_ocv1, v1[1] / len_ocv1, f / len_ocv1]
    y = [v2[0] / len_ocv2, v2[1] / len_ocv2, f / len_ocv2]
    z = cross(x, y)  # 求x和y的叉积
    R = [[x[0], y[0], z[0]], [x[1], y[1], z[1]], [x[2], y[2], z[2]]]

    p1 = [0, 0, 0]
    p2 = [width, 0, 0]
    p1_m = getmul(R, [p1])  # 旋转后的p1
    p1_m = p1_m[0]
    p2_m = getmul(R, [p2])  # 旋转后的p2
    p2_m = p2_m[0]

    p1px = line[0]
    p2px = line[1]

    i1m = [p1px[0] - oi[0], p1px[1] - oi[1]]  # 公式8

    p1_ = [i1m[0], i1m[1], f]
    p2_ = [i1m[0] + p2_m[0] - p1_m[0], i1m[1] + p2_m[1] - p1_m[1], f + p2_m[2] - p1_m[2]]  # 公式9
    q = getsec(p1_, p2_, [0, 0, 0], [p2px[0] - oi[0], p2px[1] - oi[1], f])

    len_p1p2 = getdis3(p2, p1)
    len_p1_q = getdis3(p1_, q)
    d = len_p1p2 / len_p1_q  # 公式11

    alpha = d

    T = [[p1_[0] * alpha], [p1_[1] * alpha], [p1_[2] * alpha]]

    rotation = torch.tensor(K).float().mm(torch.tensor(R).float())
    bias = torch.tensor(K).float().mm(torch.tensor(T).float())

    if getrealxy(line[1], rotation, bias)[0] < 0:
        rotation = -rotation
    return getH(rotation, bias)


def initEdaParamRng(H):
    H = torch.reshape(H, (1, 9))
    offset = torch.tensor([10]*9)
    return camParam(H, torch.add(H, offset), torch.sub(H, offset))


def calcReprojErr(camParam, uvPos, xyLens):
    fReprojErr = 0.0
    for i in range(len(xyLens)):
        oSt2dPt = uvPos[i * 2]
        oNd2dPt = uvPos[i * 2+1]
        oSt3dPt = getrealxybyH(oSt2dPt, camParam.getH().reshape(3, 3))
        oNd3dPt = getrealxybyH(oNd2dPt, camParam.getH().reshape(3, 3))
        fReprojErr += math.fabs(getdis2(oNd3dPt, oSt3dPt) - xyLens[i])
    return fReprojErr


def estEdaParamRng(pvoCamParam):
    length = len(pvoCamParam)
    tpH = torch.empty(1, 9)
    offset = torch.empty(1, 9)
    for voCam in pvoCamParam:
        tpH = torch.add(voCam.getH(), tpH)
    tpH = torch.div(tpH.float(), length)
    for voCam in pvoCamParam:
        torch.add(offset, torch.pow(torch.sub(voCam.getH(), tpH), 2))
    offset = torch.div(offset.float(), length)
    return camParam(tpH, torch.add(offset, tpH), torch.sub(tpH, offset))


def edaOpt(_H, uvPos, xyLens, outpath):
    nR = 20000
    nN = 20
    nIterNum = 10
    iIter = 0
    fReprojErrMeanPrev = 0
    EDA_REPROJ_ERR_THLD = 0.01
    sParamRng = initEdaParamRng(_H);
    if (nN >= nR):
        print("Error: Selected population should be less than initial population\n")

    voCamParam = []
    for iR in range(nR):
        voCamParam.append(sParamRng.getInsRandom())

    print("Start EDA optimization for camera calibration\n")

    while nIterNum > iIter:
        print("==== generation {}: ====\n".format(iIter))
        iProc = 0
        bProc25 = False
        bProc50 = False
        bProc75 = False
        fReprojErrMean = 0.0
        fReprojErrStd = 0.0
        for ivoCamParam in voCamParam:
            fReprojErr = calcReprojErr(ivoCamParam, uvPos, xyLens)
            ivoCamParam.setReprojErr(fReprojErr)
            fReprojErrMean += fReprojErr;
            iProc += 1

            if iProc / nR > 0.25 and bProc25 == False:
                print("25%%...")
                bProc25 = True
            if iProc / nR > 0.50 and bProc50 == False:
                print("50%%...")
                bProc50 = True
            if iProc / nR > 0.75 and bProc75 == False:
                print("75%%...")
                bProc75 = True
            fReprojErrMean /= nR;

        for ivoCamParam in voCamParam:
            fReprojErr = ivoCamParam.getReprojErr()
            fReprojErrStd += (fReprojErr - fReprojErrMean) * (fReprojErr - fReprojErrMean);

        fReprojErrStd = math.sqrt(fReprojErrStd / nR)
        print("100%%!\n");
        print("current error mean = {}\n".format(fReprojErrMean))
        print("current error standard deviation = {}\n".format(fReprojErrStd))
        for i in range(len(xyLens)):
            oSt2dPt = uvPos[i * 2]
            oNd2dPt = uvPos[i * 2 + 1]
            oSt3dPt = getrealxybyH(oSt2dPt, _H)
            oNd3dPt = getrealxybyH(oNd2dPt, _H)
            print(getdis2(oNd3dPt, oSt3dPt), xyLens[i], getdis2(oNd3dPt, oSt3dPt) - xyLens[i])

        if fReprojErrMean == 0:
            print("Camera calibration fails.\n")
            break

        if 0 < iIter and (fReprojErrMeanPrev * EDA_REPROJ_ERR_THLD > fReprojErrMeanPrev - fReprojErrMean):
            fo = open(outpath, "w+")
            out = '[' + ','.join(str(i) for i in sParamRng.getH().numpy()[0]) + ']'
            fo.write(out)
            print("Reprojection error is small enough. Stop generation.\n")
            break

        fReprojErrMeanPrev = fReprojErrMean;
        voCamParam.sort(key=lambda e: e.getReprojErr())
        voCamParam = voCamParam[:nN]
        sParamRng = estEdaParamRng(voCamParam)

        for iR in range(nR):
            voCamParam.append(sParamRng.getInsRandom())

        iIter += 1
        print("\n")

    if nIterNum <= iIter:
        print("Exit: Results can not converge.\n")
    return sParamRng.getH().reshape(3, 3)


def update(camera_id, case_id):
    config_path = osp.join('cases', 'camera_%04d' % camera_id, 'config.yaml')
    outpath = osp.join('cases', 'camera_%04d' % camera_id, 'matrix.txt')
    with open(config_path) as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    vl, vr, line, width, [w, h], uvPos, xyLens = cfg['vl'], cfg['vr'], cfg['line'],\
                                                 cfg['width'], cfg['size'], cfg['uvPos'], cfg['xyLens']
    _H = get_parameters_auto(vl, vr, line, width, w, h)
    _H = edaOpt(_H, uvPos, xyLens, outpath)
    for i in range(len(xyLens)):
        oSt2dPt = uvPos[i * 2]
        oNd2dPt = uvPos[i * 2+1]
        oSt3dPt = getrealxybyH(oSt2dPt, _H)
        oNd3dPt = getrealxybyH(oNd2dPt, _H)
        print(getdis2(oNd3dPt, oSt3dPt), xyLens[i], getdis2(oNd3dPt, oSt3dPt) - xyLens[i])


if __name__ == '__main__':
    update(0, 0)
