'''
Author: caishuyang
Date: 2023-01-08 16:52:46
LastEditors: caishuyang
LastEditTime: 2023-01-09 19:11:37
Description: 
'''

from Ball import Ball
from Contain import Contain
import numpy as np
import yaml
import math
import cv2


class Generator():
    def __init__(self):

        print('-----开始初始化-----')
        self.cont = Contain()
        self.balldic, self.cdic, self.fdic, self.deltat = self.getDic()
        self.deltat = float(self.deltat)
        self.BalNalist = self.getBallName()
        self.maxR, self.minR = self.getMaxMin()
        self.balls = self.getBalls()
        print('-----初始化结束-----')
        print('------模拟开始------')

    def getDic(self):
        yamlPath = 'config.yaml'
        flag = True
        try:
            f = open(yamlPath, 'r', encoding='utf-8')
        except:
            print('无法打开', yamlPath)
            flag = False
        if flag:
            cont = f.read()
            param = yaml.load(cont, Loader=yaml.FullLoader)
            balldic = param['Ball']
            fdic = param['Friction']
            cdic = param['Collision']
            t = param['DeltaTime']
            return balldic, cdic, fdic, t

    def getBallName(self):
        l = []
        for key, value in enumerate(self.balldic):
            l.append((value, int(self.balldic[value]['num'])))
        return l

    def getMaxMin(self):
        Rmax = float(self.balldic[self.BalNalist[0][0]]['r'])
        Rmin = float(self.balldic[self.BalNalist[0][0]]['r'])
        for val in self.BalNalist:
            if Rmax < float(self.balldic[val[0]]['r']):
                Rmax = float(self.balldic[val[0]]['r'])
            if Rmin > float(self.balldic[val[0]]['r']):
                Rmin = float(self.balldic[val[0]]['r'])
        return Rmax, Rmin

    def getBalls(self):
        count = 0
        l = []
        xy = []
        jg = 1.1*self.maxR
        for y in range(int(self.cont.R//jg)):
            truey = self.cont.R-jg-2*y*jg
            halflong = ((self.cont.R-jg)**2-truey**2)**0.5
            xy.append((0, truey))
            for x in range(int(halflong//(2*jg))):
                xy.append((2*(x+1)*jg, truey))
                xy.append((-2*(x+1)*jg, truey))
        countmax = len(xy)

        for val in self.BalNalist:
            for i in range(val[1]):
                tempbal = Ball(val[0])
                tempbal.x = xy[count][0]
                tempbal.y = xy[count][1]
                tempbal.vx = -self.cont.vx
                tempbal.vy = -self.cont.vy

                l.append(tempbal)
                count = count+1
                if count == countmax:
                    break
        return l

    def getEc(self, name1, name2):
        ec = self.cdic[str(name1)][str(name2)]
        return ec

    def getFr(self, name1, name2):
        f = self.fdic[str(name1)][str(name2)]
        return f

    def updateA(self):
        g = 9.8
        for i in range(len(self.balls)):
            miu = self.getFr(self.cont, self.balls[i])
            ag = self.cont.omiga**2*self.cont.rotR
            f_m = miu*g
            agx = ag * math.cos(self.cont.omiga*self.cont.t)
            agy = ag * math.sin(self.cont.omiga*self.cont.t)

            if self.balls[i].vx**2+self.balls[i].vy**2 < (self.minR/100)**2:
                if ag <= f_m:
                    f_m = ag

                afx = -f_m*math.cos(self.cont.omiga*self.cont.t)
                afy = -f_m*math.sin(self.cont.omiga*self.cont.t)

            else:
                afx = f_m*(-self.balls[i].vx) / \
                    (self.balls[i].vx**2+self.balls[i].vy**2)**0.5
                afy = f_m*(-self.balls[i].vy) / \
                    (self.balls[i].vx**2+self.balls[i].vy**2)**0.5

            self.balls[i].ax = afx+agx
            self.balls[i].ay = afy+agy

    def updateVz(self):
        for i in range(len(self.balls)):
            self.balls[i].vx += self.deltat*self.balls[i].ax
            self.balls[i].vy += self.deltat*self.balls[i].ay

    def updateX(self):
        flag = []
        for i in range(len(self.balls)):
            test = True
            newx = self.balls[i].x+self.deltat*self.balls[i].vx
            newy = self.balls[i].y+self.deltat*self.balls[i].vy
            if (newx**2+newy**2)**0.5+self.balls[i].r/2 > self.cont.R:
                test = False
                flag.append(test)
                continue
            for j in range(len(self.balls)):
                if j != i:
                    newx2 = self.balls[j].x+self.deltat*self.balls[j].vx
                    newy2 = self.balls[j].y+self.deltat*self.balls[j].vy
                    if ((newx-newx2)**2+(newy-newy2)**2)**0.5 < 0.75*(self.balls[j].r+self.balls[i].r):
                        test = False
                        break
            flag.append(test)
        try:
            for i in range(len(self.balls)):
                if flag[i]:
                    self.balls[i].x += self.deltat*self.balls[i].vx
                    self.balls[i].y += self.deltat*self.balls[i].vy
                else:
                    self.balls[i].vx *= 0
                    self.balls[i].vy *= 0
        except:
            print(len(flag))

        self.cont.t += self.deltat
        self.cont.getAll()

    def collide(self, object1, object2):
        ec = self.getEc(object1, object2)
        flag = 0
        if str(object1) == 'Contain':
            vx0 = object2.vx
            vy0 = object2.vy
            x = object2.x
            y = object2.y
            cos = x/(x**2+y**2)**0.5
            sin = y/(x**2+y**2)**0.5
            I_m = (1+ec)*(vx0*cos+vy0*sin)
            if I_m <= 0.0001:
                I_m = 0
                flag = 1
            object2.vx = vx0-I_m*cos
            object2.vy = vy0-I_m*sin

        elif str(object2) == 'Contain':
            vx0 = object1.vx
            vy0 = object1.vy
            x = object1.x
            y = object1.y
            cos = x/(x**2+y**2)**0.5
            sin = y/(x**2+y**2)**0.5

            I_m = (1+ec)*(vx0*cos+vy0*sin)
            if I_m <= 0.0001:
                I_m = 0
                flag = 1
            object1.vx = vx0-I_m*cos
            object1.vy = vy0-I_m*sin

        else:
            v1x = object1.vx
            v1y = object1.vy
            v2x = object2.vx
            v2y = object2.vy
            m1 = object1.m
            m2 = object2.m
            x = object2.x-object1.x
            y = object2.y-object1.y
            cos = x/(x**2+y**2)**0.5
            sin = y/(x**2+y**2)**0.5
            It = (1+ec)*((v1x-v2x)*cos+(v1y-v2y)*sin)/(1/m1+1/m2)
            if It < 0.0001:
                It = 0
                flag = 1
            object1.vx = v1x-It/m1*cos
            object1.vy = v1x-It/m1*sin
            object2.vx = v2x+It/m2*cos
            object2.vy = v2y+It/m2*sin
        return flag

    def isnear(self, object1, object2):
        flag = False
        if str(object1) == 'Contain':
            r = (object2.x**2+object2.y**2)**0.5
            if object1.R-object2.r-r < self.minR/1000:
                flag = True

        elif str(object2) == 'Contain':
            r = (object1.x**2+object1.y**2)**0.5
            if object2.R-object1.r-r < self.minR/1000:
                flag = True
        else:
            r = ((object2.x-object1.x)**2+(object2.y-object1.y)**2)**0.5
            if r-object1.r-object2.r < self.minR/100:
                flag = True

        return flag

    def cProgress(self):

        for i in range(len(self.balls)):
            near = 0
            if self.isnear(self.balls[i], self.cont):
                near += 1
                self.collide(self.balls[i], self.cont)

            for j in range(len(self.balls)):
                if near >= 6:
                    break
                if j != i:
                    if self.isnear(self.balls[i], self.balls[j]):
                        near += 1
                        self.collide(self.balls[i], self.balls[j])

    def painter2D(self):
        canvas = 255*np.ones((800, 800, 3), dtype=np.uint8)
        bl = 300/(self.cont.R+self.cont.rotR)
        R0 = int(bl*self.cont.R)
        x0 = int(bl*self.cont.x+400)
        y0 = int(-bl*self.cont.y+400)
        cv2.circle(canvas, (x0, y0), R0, (0, 0, 0))
        for ball in self.balls:
            r = int(bl*ball.r)
            x = int(bl*(ball.x+self.cont.x)+400)
            y = int(-bl*(ball.y+self.cont.y)+400)
            try:

                if self.isOp(ball):
                    cv2.circle(canvas, (int(x), int(y)), r, (0, 0, 255))
                else:
                    cv2.circle(canvas, (int(x), int(y)), r, (255, 0, 0))
            except:
                print((x, y))
        txt = 'The total energy of ' + \
            str(len(self.balls)) + ' balls: ' + \
            str(format(1000*self.calE(), '.3f'))+'e-3J'
        cv2.putText(canvas, txt, (5, 50), cv2.FONT_HERSHEY_SIMPLEX,
                    0.75, (0, 127, 255), 2)
        cv2.imshow('show', canvas)

    def isOp(self, ball):
        flag = False
        vxt = ball.vx+self.cont.vx
        vyt = ball.vy+self.cont.vy
        if vxt*self.cont.vx+vyt*self.cont.vy < 0:
            flag = True
        return flag

    def calE(self):
        E = 0
        for ball in self.balls:
            E += 0.5*ball.m*((ball.vx)**2+(ball.vy)**2)
        return E