using UnityEngine;

public class AntColonyOpt : MonoBehaviour
{
    // min angle that the min path can have
    public float angleOffset;
    public float crossTweakForce;
    public float crossDiv;
    public bool bParticle;
    public bool bDrawLine;
    public float emphasizeBest;
    public int numInGroup;
    public int preTimes;
    public float disPow;
    public float weightPow;
    public float dropPow;
    public float deltaPhe;
    public int pheMaxTimes;
    public float pheDrop;
    public int numPoint;
    public float xRange;
    public float yRange;
    public float minDis;
    public GameObject pointPrefab;
    public GameObject particalPrefab;
    [HideInInspector]
    public int preTimeLeft;
    [HideInInspector]
    public Vector2[] pPos;
    [HideInInspector]
    public float[] dis;
    [HideInInspector]
    public float[] pheromone;
    [HideInInspector]
    public float pheMax;
    [HideInInspector]
    public int numPath;
    [HideInInspector]
    public float minPathDis;
    [HideInInspector]
    public int[] minPath;
    [HideInInspector]
    public bool bSpace;
    [HideInInspector]
    public bool bAuto;
    [HideInInspector]
    public bool bFirstRound;
    [HideInInspector]
    public float pheInit;
    [HideInInspector]
    public float pheMin;
    [HideInInspector]
    public NowPath nowPath;
    [HideInInspector]
    public int minCrossCnt;
    [HideInInspector]
    public CrossPath crossPath;
    [HideInInspector]
    public LineRenderer lineRenderer;
    [HideInInspector]
    public static Particle[] particles;
    [HideInInspector]
    public int turnCnt;

    int C2(int n)
    {
        int result = 0;
        for (int i = n - 1; i > 0; i --) result += i;
        return result;
    }

    // passing the index of two point
    // return the index of the path connecting them
    int GetPathIndex(int cur1, int cur2)
    {
        if (cur1 > cur2)
        {
            int tmp = cur1;
            cur1 = cur2;
            cur2 = tmp;
        }
        int index = 0;
        for (int i = 1; i <= cur1; i++) index += numPoint - i;
        index += cur2 - cur1 - 1;
        return index;
    }

    void Awake()
    {
        nowPath = FindObjectOfType<NowPath>();
        crossPath = FindObjectOfType<CrossPath>();
        //numPath = (int)C(numPoint, 2);
        numPath = C2(numPoint);
        for (int i = 0; i < numPath && bParticle; i ++)
            Instantiate(particalPrefab);
    }

    void Start()
    {
        turnCnt = 0;
        preTimeLeft = 0;
        bAuto = false;
        bFirstRound = true;
        pPos = new Vector2[numPoint];
        minPath = new int[numPoint];
        float sqrMinDis = minDis * minDis;

        // gen random points
        for (int i = 0; i < numPoint; i++)
        {
            minPath[i] = i;

            Vector2 pos = new Vector2(Random.Range(-xRange, xRange), Random.Range(-yRange, yRange));
            bool tooClose = false;
            for (int j = 0; j < i; j++)
                if ((pos - pPos[j]).sqrMagnitude < sqrMinDis)
                {
                    i--;
                    tooClose = true;
                    break;
                }
            if (tooClose) continue;
            GameObject point = Instantiate(pointPrefab);
            point.transform.position = pos;
            point.name = i.ToString();
            pPos[i] = pos;
        }

        // calculate pheromone parameters
        pheMax = deltaPhe * pheMaxTimes * numInGroup;
        pheInit = pheMax / 5;
        pheMin = pheInit * Mathf.Pow(pheDrop, 5);

        // initialization of arrays
        pheromone = new float[numPath];
        for (int i = 0; i < numPath; i++)
            pheromone[i] = pheInit;
        dis = new float[numPath];
        particles = FindObjectsOfType<Particle>();

        // calculate distance
        // set rotation and length of particles
        for (int i = 0; i < numPoint - 1; i++)
            for (int j = i + 1; j < numPoint; j++)
            {
                int nowIndex = GetPathIndex(i, j);
                dis[nowIndex] = (pPos[j] - pPos[i]).magnitude;

                if(bParticle)
                {
                    particles[nowIndex].transform.position = pPos[i];

                    float angle = Vector2.Angle(Vector2.up, pPos[j] - pPos[i]);
                    if (Vector3.Cross(Vector2.up, pPos[j] - pPos[i]).z > 0) angle = -angle;
                    particles[nowIndex].SetParticleRot(angle);

                    particles[nowIndex].SetParticleDis(dis[nowIndex]);
                }
            }

        // initialize line renderer
        lineRenderer = GetComponent<LineRenderer>();
        lineRenderer.positionCount = numPoint + 1;
    }

    void Update()
    {
        bSpace = false;

        if (Input.GetKeyDown(KeyCode.W))
        {
            preTimeLeft = preTimes;
        }
        if (preTimeLeft > 0)
        {
            preTimeLeft--;
            Opt();
            if (preTimeLeft == 0) print("preTimeFinished");
        }
        if (Input.GetKeyDown(KeyCode.Space))
        {
            bSpace = true;
            Opt();
        }
        if (Input.GetKeyDown(KeyCode.S))
        {
            GameObject[] points = GameObject.FindGameObjectsWithTag("Point");
            for (int i = 0; i < numPoint; i++)
            {
                Destroy(points[i]);
                nowPath.ClearPath();
            }
            Start();
        }
        if (Input.GetKeyDown(KeyCode.A))
        {
            print("AUTO");
            bAuto = true;
        }
        if (Input.GetKeyDown(KeyCode.D))
        {
            print("MANUAL");
            bAuto = false;
            print(turnCnt);
        }
        if (bAuto) Opt();
    }

    void Opt()
    {
        turnCnt++;
        int[] path = new int[numPoint * numInGroup];
        for (int antCnt = 0; antCnt < numInGroup; antCnt++)
        {
            // offset of starting point index
            int offset = Random.Range(0, numPoint);
            int nowP = offset;
            int numVisited = 0;
            bool[] visited = new bool[numPoint];
            visited[offset] = true;
            path[numPoint * antCnt] = offset;

            for (int turn = 0; turn < numPoint - 1; turn++)
            {
                float[] tempDis = new float[numPoint];
                float[] tempPhe = new float[numPoint];
                float[] weight = new float[numPoint];

                for (int i = 0; i < numPoint; i++)
                    if (!visited[i])
                    {
                        int nowIndex = GetPathIndex(nowP, i);
                        float nowDis = Mathf.Pow(1f / dis[nowIndex], disPow);
                        // weight will be diveded by total then
                        tempPhe[i] = pheromone[nowIndex];
                        tempDis[i] = nowDis;
                    }

                // calculate weight
                float totalWeight = 0;
                for (int i = 0; i < numPoint; i++)
                    if (!visited[i])
                    {
                        weight[i] = Mathf.Pow(tempDis[i] * tempPhe[i], weightPow);
                        // reduce chance of getting cross
                        for (int j = 0; j < numVisited - 1; j ++)
                            if(DetectCross(pPos[i], pPos[nowP], pPos[path[j + numPoint * antCnt]], pPos[path[j + 1 + numPoint * antCnt]]))
                                weight[i] /= crossDiv;

                        totalWeight += weight[i];
                    }
                for (int i = 0; i < numPoint; i++)
                    if (!visited[i])
                        weight[i] /= totalWeight;

                // Random.value returns in range [0,1]
                float choice = Random.value;

                int decideIndex = offset;
                int maxCnt = 0;
                while (maxCnt < 2 * numPoint)
                {
                    int judgingIndex = decideIndex % numPoint;
                    choice -= weight[judgingIndex];
                    if (choice <= 0 && !visited[judgingIndex])
                    {
                        numVisited++;
                        path[numVisited + numPoint * antCnt] = judgingIndex;
                        nowP = judgingIndex;
                        visited[judgingIndex] = true;

                        break;
                    }
                    decideIndex++;
                    maxCnt++;
                }
            }

            // update when a turn is ending
            if (antCnt == numInGroup - 1)
            {
                float localMinPathDis = float.MaxValue;
                float[] disGroup = new float[numInGroup];
                int localMinIndex = 0;

                // get localMinPathDis and its index
                for (int antIndex = 0; antIndex < numInGroup; antIndex++)
                {
                    for (int i = 0; i < numPoint; i++)
                        disGroup[antIndex] += dis[GetPathIndex(path[i + numPoint * antIndex],
                            (i + 1) % numPoint == 0 ? path[numPoint * antIndex] : path[i + 1 + numPoint * antIndex])];
                    if (disGroup[antIndex] < localMinPathDis)
                    {
                        localMinPathDis = disGroup[antIndex];
                        localMinIndex = antIndex;
                    }
                }
                
                int[] localMinPath = new int[numPoint];
                for (int i = 0; i < numPoint; i++)
                    localMinPath[i] = path[i + localMinIndex * numPoint];

                if(bFirstRound)
                {
                    minPathDis = localMinPathDis;
                    minCrossCnt = CountCross(localMinPath);
                    for (int i = 0; i < numPoint; i++)
                    {
                        if(bDrawLine) lineRenderer.SetPosition(i, pPos[localMinPath[i]]);
                        minPath[i] = localMinPath[i];
                    }
                    if(bDrawLine) lineRenderer.SetPosition(numPoint, pPos[localMinPath[0]]);
                }

                // pheromone drop
                for (int i = 0; i < numPath; i++)
                {
                    pheromone[i] *= pheDrop;
                    if (pheromone[i] < pheMin) pheromone[i] = pheMin;
                }

                // add pheromone of the turn
                for (int antIndex = 0; antIndex < numInGroup; antIndex++)
                    for (int i = 0; i < numPoint; i++)
                    {
                        int nowIndex = GetPathIndex(path[i + numPoint * antIndex], (i + 1) % numPoint == 0
                            ? path[numPoint * antIndex] : path[i + 1 + numPoint * antIndex]);
                        // shorter path impacts pheromone more
                        pheromone[nowIndex] += deltaPhe * Mathf.Pow(localMinPathDis / disGroup[antIndex], dropPow);
                        if (pheromone[nowIndex] > pheMax) pheromone[nowIndex] = pheMax;
                    }

                // emphasize best solution but avoid its crossing parts and small angle
                for (int i = 0; i < numPoint; i++)
                {
                    int cur1 = minPath[i];
                    int cur2 = (i + 1) % numPoint == 0 ? minPath[0] : minPath[i + 1];
                    int nowIndex = GetPathIndex(cur1, cur2);
                    pheromone[nowIndex] += deltaPhe * emphasizeBest;
                    if (pheromone[nowIndex] > pheMax) pheromone[nowIndex] = pheMax;

                    // avoid crossing
                    for (int j = 0; j < i - 1; j++)
                        if (DetectCross(pPos[cur1], pPos[cur2], pPos[minPath[j]], pPos[minPath[j + 1]]))
                        {
                            int anoIndex = GetPathIndex(minPath[j], minPath[j + 1]);
                            pheromone[nowIndex] = pheMin;
                            pheromone[anoIndex] = pheMin;

                            pheromone[GetPathIndex(cur1, minPath[j])] += deltaPhe * crossTweakForce;
                            if (pheromone[GetPathIndex(cur1, minPath[j])] > pheMax) pheromone[nowIndex] = pheMax;
                            pheromone[GetPathIndex(cur2, minPath[j + 1])] += deltaPhe * crossTweakForce;
                            if (pheromone[GetPathIndex(cur2, minPath[j + 1])] > pheMax) pheromone[nowIndex] = pheMax;
                        }
                    // avoid small angle
                    if (i != 0 && DetectSmallAngle(pPos[cur1] - pPos[cur2], pPos[cur1] - pPos[minPath[i - 1]], angleOffset))
                    {
                        int anoIndex = GetPathIndex(minPath[i - 1], cur1);

                        pheromone[nowIndex] -= deltaPhe * emphasizeBest;
                        pheromone[anoIndex] -= deltaPhe * emphasizeBest;
                    }
                }

                // set particle size based on pheromone
                for (int i = 0; i < numPoint && bParticle; i++)
                {
                    int mostPheIndex = i + 1;
                    for (int j = i + 1; j < numPoint; j++)
                    {
                        int nowIndex = GetPathIndex(i, j);
                        float sizeRatio = (pheromone[nowIndex] - pheMin) / (pheMax - pheMin);
                        particles[nowIndex].SetParticleSize(sizeRatio);
                    }
                }
                    
                // update minPathDis and display
                if (localMinPathDis < minPathDis)
                {
                    minCrossCnt = CountCross(localMinPath);
                    minPathDis = localMinPathDis;
                    print(minPathDis);
                    print(minCrossCnt);
                    for (int i = 0; i < numPoint; i++)
                    {
                        if(bDrawLine) lineRenderer.SetPosition(i, pPos[localMinPath[i]]);
                        minPath[i] = localMinPath[i];
                    }
                if(bDrawLine) lineRenderer.SetPosition(numPoint, pPos[localMinPath[0]]);
                }

                // show Recent Path
                if (bSpace && !bParticle) nowPath.DrawNowPath(path, pPos, numPoint, localMinIndex);

                bFirstRound = false;
            }
        }
    }

    int CountCross(int[] path)
    {
        int crossCnt = 0;
        for (int i = 2; i < path.Length; i++)
        {
            Vector3 pA = pPos[path[i]];
            Vector3 pB = pPos[i + 1 == path.Length ? path[0] : path[i + 1]];
            for (int j = 0; j < i - 1; j++)
            {
                Vector3 pC = pPos[path[j]];
                Vector3 pD = pPos[path[j + 1]];
                if(DetectCross(pA, pB, pC, pD)) crossCnt++;
            }
        }
        return crossCnt;
    }

    bool DetectCross(Vector3 pA, Vector3 pB, Vector3 pC, Vector3 pD)
    {
        if (Vector3.Dot(Vector3.Cross(pA - pD, pC - pD), Vector3.Cross(pB - pD, pC - pD)) < 0
            && Vector3.Dot(Vector3.Cross(pC - pB, pA - pB), Vector3.Cross(pD - pB, pA - pB)) < 0)
        {
            return true;;
        }
        return false;
    }

    bool DetectSmallAngle(Vector2 vec1, Vector2 vec2, float angleOffset)
    {
        if (Vector2.Angle(vec1, vec2) < angleOffset)
            return true;
        return false;
    }
}

using UnityEngine;

public class Particle : MonoBehaviour
{
    public float sizePow;
    [HideInInspector]
    public new ParticleSystem particleSystem;
    void Awake()
    {
        particleSystem = GetComponent<ParticleSystem>();
    }
    public void SetParticleRot(float angle)
    {
        var sh = particleSystem.shape;
        sh.enabled = true;
        sh.shapeType = ParticleSystemShapeType.Rectangle;
        sh.rotation = new Vector3(0f, angle, 0f);
    }
    public void SetParticleDis(float dis)
    {
        var main = particleSystem.main;
        float speed = main.startSpeed.constant;
        main.startLifetime = dis / speed;
    }
    // max Size = 2.5, min Size = 0.5
    public void SetParticleSize(float sizeRatio)
    {
        var main = particleSystem.main;
        sizeRatio = Mathf.Pow(sizeRatio, sizePow);
        float size = sizeRatio * 2.5f;
        main.startSize = size > 0.5f? size : 0f;
    }
}