using System.Collections.Generic;
using UnityEngine;

namespace FishROV.AvoidanceBenchmark
{
    public sealed class SamplingLocalAvoidanceAdapter : IAvoidanceAdapter
    {
        private const float TimeHorizon = 1.25f;
        private const float NeighbourRangeMultiplier = 5.5f;
        private const float GoalWeight = 0.85f;
        private const float SmoothWeight = 0.12f;
        private const float CollisionWeight = 18f;
        private const float CurrentOverlapWeight = 80f;
        private const float BoundaryWeight = 12f;
        private const int DirectionSamples = 16;

        private readonly Dictionary<BenchmarkAgent, AgentState> states =
            new Dictionary<BenchmarkAgent, AgentState>(512);

        private readonly List<AgentState> orderedStates = new List<AgentState>(512);
        private readonly Dictionary<Vector2Int, List<AgentState>> buckets =
            new Dictionary<Vector2Int, List<AgentState>>(1024);

        private AvoidanceBenchmarkConfig config;
        private float cellSize;

        public string Name => "采样避障（学习版）";
        public bool IsAvailable => true;
        public string Status => "学习版：空间哈希找邻居，采样候选速度，用预测碰撞/平滑/目标偏差/边界惩罚打分。";

        public void Initialize(AvoidanceBenchmarkConfig config)
        {
            this.config = config.Clone();
            states.Clear();
            orderedStates.Clear();
            buckets.Clear();
            cellSize = config.FishRadius * NeighbourRangeMultiplier;
        }

        public BenchmarkAgent CreateAgent(BenchmarkAgentSpawnInfo spawnInfo, Transform parent)
        {
            var agent = BenchmarkAgent.CreatePrimitive(spawnInfo, parent);
            var state = new AgentState(agent);
            states[agent] = state;
            orderedStates.Add(state);
            return agent;
        }

        public void SetPreferredVelocity(BenchmarkAgent agent, Vector3 preferredVelocity)
        {
            if (!states.TryGetValue(agent, out var state))
            {
                agent.ApplyVelocity(preferredVelocity);
                return;
            }

            state.PreferredVelocity = ClampForMode(preferredVelocity, agent.MaxSpeed);
        }

        public void SetTarget(BenchmarkAgent agent, Vector3 target)
        {
            agent.Target = target;
        }

        public void Tick(float deltaTime)
        {
            if (deltaTime <= 0f)
            {
                return;
            }

            BuildBuckets();

            for (var i = 0; i < orderedStates.Count; i++)
            {
                var state = orderedStates[i];
                if (state.Agent == null)
                {
                    continue;
                }

                var velocity = ChooseVelocity(state, deltaTime);
                state.Agent.ApplyVelocity(velocity);
                state.LastVelocity = velocity;
            }
        }

        public void Dispose()
        {
            states.Clear();
            orderedStates.Clear();
            buckets.Clear();
        }

        private Vector3 ChooseVelocity(AgentState state, float deltaTime)
        {
            var preferred = state.PreferredVelocity;
            var bestVelocity = preferred;
            var bestScore = EvaluateVelocity(state, preferred, deltaTime);

            TestCandidate(state, Vector3.zero, deltaTime, ref bestVelocity, ref bestScore);

            var flatPreferred = Flatten(preferred);
            var baseDirection = flatPreferred.sqrMagnitude > 0.0001f
                ? flatPreferred.normalized
                : ForwardFromAgent(state.Agent);

            var maxSpeed = state.Agent.MaxSpeed;
            for (var i = 0; i < DirectionSamples; i++)
            {
                var angle = i * 360f / DirectionSamples;
                var direction = Quaternion.AngleAxis(angle, Vector3.up) * baseDirection;
                TestCandidate(state, RestoreVertical(direction * maxSpeed, preferred), deltaTime, ref bestVelocity, ref bestScore);
                TestCandidate(state, RestoreVertical(direction * maxSpeed * 0.55f, preferred), deltaTime, ref bestVelocity, ref bestScore);
            }

            return ClampForMode(bestVelocity, maxSpeed);
        }

        private void TestCandidate(
            AgentState state,
            Vector3 candidate,
            float deltaTime,
            ref Vector3 bestVelocity,
            ref float bestScore)
        {
            var score = EvaluateVelocity(state, candidate, deltaTime);
            if (score < bestScore)
            {
                bestScore = score;
                bestVelocity = candidate;
            }
        }

        private float EvaluateVelocity(AgentState state, Vector3 candidate, float deltaTime)
        {
            candidate = ClampForMode(candidate, state.Agent.MaxSpeed);
            var preferred = state.PreferredVelocity;
            var score = GoalWeight * (candidate - preferred).sqrMagnitude;
            score += SmoothWeight * (candidate - state.LastVelocity).sqrMagnitude;
            score += BoundaryPenalty(state.Agent.transform.position, candidate, deltaTime);

            var position = state.Agent.transform.position;
            var centerCell = ToCell(position);
            var queryRange = Mathf.CeilToInt(
                state.Agent.Radius * NeighbourRangeMultiplier / cellSize);
            queryRange = Mathf.Clamp(queryRange, 1, 4);

            for (var x = -queryRange; x <= queryRange; x++)
            {
                for (var z = -queryRange; z <= queryRange; z++)
                {
                    var cell = new Vector2Int(centerCell.x + x, centerCell.y + z);
                    if (!buckets.TryGetValue(cell, out var bucket))
                    {
                        continue;
                    }

                    for (var i = 0; i < bucket.Count; i++)
                    {
                        var other = bucket[i];
                        if (other == state || other.Agent == null)
                        {
                            continue;
                        }

                        score += CollisionPenalty(state, other, candidate);
                    }
                }
            }

            return score;
        }

        private float CollisionPenalty(AgentState state, AgentState other, Vector3 candidate)
        {
            var position = state.Agent.transform.position;
            var otherPosition = other.Agent.transform.position;
            var relativePosition = otherPosition - position;
            var combinedRadius = state.Agent.Radius + other.Agent.Radius;
            var neighbourRange = Mathf.Max(combinedRadius * 2f, state.Agent.Radius * NeighbourRangeMultiplier);

            if (config.DimensionMode != AvoidanceDimensionMode.XYZ3D)
            {
                relativePosition.y = 0f;
            }

            var distanceSq = relativePosition.sqrMagnitude;
            if (distanceSq > neighbourRange * neighbourRange)
            {
                return 0f;
            }

            var relativeVelocity = candidate - other.LastVelocity;
            if (config.DimensionMode != AvoidanceDimensionMode.XYZ3D)
            {
                relativeVelocity.y = 0f;
            }

            var clearanceRadius = combinedRadius * 1.08f;
            var currentDistance = Mathf.Sqrt(Mathf.Max(0.0001f, distanceSq));
            var currentClearance = currentDistance - clearanceRadius;
            var penalty = 0f;

            if (currentClearance < 0f)
            {
                var overlap = -currentClearance / clearanceRadius;
                penalty += CurrentOverlapWeight * overlap * overlap;
            }

            var relativeSpeedSq = relativeVelocity.sqrMagnitude;
            if (relativeSpeedSq < 0.0001f)
            {
                return penalty;
            }

            var closestTime = Mathf.Clamp(
                -Vector3.Dot(relativePosition, relativeVelocity) / relativeSpeedSq,
                0f,
                TimeHorizon);
            var closest = relativePosition + relativeVelocity * closestTime;
            var closestDistance = closest.magnitude;
            var predictedClearance = closestDistance - clearanceRadius;
            if (predictedClearance < 0f)
            {
                var risk = -predictedClearance / clearanceRadius;
                var urgency = 1f + (TimeHorizon - closestTime) / TimeHorizon;
                penalty += CollisionWeight * risk * risk * urgency;
            }

            return penalty;
        }

        private float BoundaryPenalty(Vector3 position, Vector3 velocity, float deltaTime)
        {
            var next = position + velocity * deltaTime * TimeHorizon;
            var flat = new Vector2(next.x, next.z);
            var softRadius = config.ArenaRadius * 0.92f;
            if (flat.magnitude <= softRadius)
            {
                return 0f;
            }

            var overflow = (flat.magnitude - softRadius) / Mathf.Max(0.001f, config.ArenaRadius - softRadius);
            return BoundaryWeight * overflow * overflow;
        }

        private void BuildBuckets()
        {
            buckets.Clear();
            for (var i = 0; i < orderedStates.Count; i++)
            {
                var state = orderedStates[i];
                if (state.Agent == null)
                {
                    continue;
                }

                var cell = ToCell(state.Agent.transform.position);
                if (!buckets.TryGetValue(cell, out var bucket))
                {
                    bucket = new List<AgentState>(8);
                    buckets[cell] = bucket;
                }

                bucket.Add(state);
            }
        }

        private Vector2Int ToCell(Vector3 position)
        {
            return new Vector2Int(
                Mathf.FloorToInt(position.x / cellSize),
                Mathf.FloorToInt(position.z / cellSize));
        }

        private Vector3 ClampForMode(Vector3 velocity, float maxSpeed)
        {
            if (config.DimensionMode != AvoidanceDimensionMode.XYZ3D)
            {
                velocity.y = 0f;
            }

            return Vector3.ClampMagnitude(velocity, maxSpeed);
        }

        private Vector3 RestoreVertical(Vector3 flatVelocity, Vector3 preferred)
        {
            if (config.DimensionMode == AvoidanceDimensionMode.XYZ3D)
            {
                flatVelocity.y = preferred.y;
            }
            else
            {
                flatVelocity.y = 0f;
            }

            return flatVelocity;
        }

        private static Vector3 Flatten(Vector3 value)
        {
            value.y = 0f;
            return value;
        }

        private static Vector3 ForwardFromAgent(BenchmarkAgent agent)
        {
            var forward = agent != null ? agent.transform.forward : Vector3.forward;
            forward.y = 0f;
            return forward.sqrMagnitude > 0.0001f ? forward.normalized : Vector3.forward;
        }

        private sealed class AgentState
        {
            public AgentState(BenchmarkAgent agent)
            {
                Agent = agent;
            }

            public BenchmarkAgent Agent { get; }
            public Vector3 PreferredVelocity { get; set; }
            public Vector3 LastVelocity { get; set; }
        }
    }
}
