#include <Arduino.h>
#include <cmath>
#include <iostream>
#include <vector>
#include <unordered_map>
#include <queue>
#include <limits>
#include <algorithm>

using namespace std;

enum Direction { NORTH, EAST, SOUTH, WEST };
Direction currentDirection;
int currentLocation;

struct Edge {
    int neighborId;
    float cost;
};

struct Node {
    int id;
    float x, y;
    string label; // "ROOM", "STATION", "INTERSECTION"
    vector<Edge> neighbors;

    // A* metadata
    float gCost = numeric_limits<float>::infinity();
    float hCost = 0;
    float fCost = 0;
    int parentId = -1;

    void reset() {
        gCost = numeric_limits<float>::infinity();
        hCost = 0;
        fCost = 0;
        parentId = -1;
    }
};

vector<Node> allNodes;


// ====== THUẬT TOÁN A* ======

struct CompareF {
    bool operator()(const pair<int, float>& a, const pair<int, float>& b) {
        return a.second > b.second; // min-heap dựa vào fCost
    }
};

float heuristic(const Node& a, const Node& b) {
    // Dùng khoảng cách Euclidean hoặc Manhattan
    return abs(a.x - b.x) + abs(a.y - b.y); // Manhattan
    // return sqrt((a.x - b.x)*(a.x - b.x) + (a.y - b.y)*(a.y - b.y)); // Euclidean
}

void resetGraph(vector<Node>& graph) {
    for (auto& node : graph)
        node.reset();
}

vector<int> aStar(int startId, int goalId, vector<Node>& graph) {
    resetGraph(graph);

    graph[startId].gCost = 0;
    graph[startId].hCost = heuristic(graph[startId], graph[goalId]);
    graph[startId].fCost = graph[startId].hCost;

    priority_queue<pair<int, float>, vector<pair<int, float>>, CompareF> openSet;
    unordered_map<int, bool> inOpenSet;

    openSet.push({startId, graph[startId].fCost});
    inOpenSet[startId] = true;

    while (!openSet.empty()) {
        int currentId = openSet.top().first;
        openSet.pop();
        inOpenSet[currentId] = false;

        if (currentId == goalId) {
            // reconstruct path (danh sách các id)
            vector<int> path;
            for (int id = goalId; id != -1; id = graph[id].parentId)
                path.push_back(id);
            reverse(path.begin(), path.end());
            return path;
        }

        for (const Edge& edge : graph[currentId].neighbors) {
            int neighborId = edge.neighborId;
            float tentativeG = graph[currentId].gCost + edge.cost;

            if (tentativeG < graph[neighborId].gCost) {
                graph[neighborId].gCost = tentativeG;
                graph[neighborId].hCost = heuristic(graph[neighborId], graph[goalId]);
                graph[neighborId].fCost = graph[neighborId].gCost + graph[neighborId].hCost;
                graph[neighborId].parentId = currentId;

                if (!inOpenSet[neighborId]) {
                    openSet.push({neighborId, graph[neighborId].fCost});
                    inOpenSet[neighborId] = true;
                }
            }
        }
    }

    return {}; // Không tìm thấy đường
}


// ====== HÀM TÌM TRẠM GẦN NHẤT ======

int findNearestStation(int fromId, vector<Node>& graph) {
    float minCost = numeric_limits<float>::infinity();
    int nearestStationId = -1;

    for (const Node& node : graph) {
        if (node.label == "STATION") {
            vector<int> path = aStar(fromId, node.id, graph);
            if (!path.empty()) {
                float cost = graph[node.id].gCost;
                if (cost < minCost) {
                    minCost = cost;
                    nearestStationId = node.id;
                }
            }
        }
    }

    return nearestStationId;
}


int motor1pin1 = 25;
int motor1pin2 = 26;
int motor2pin1 = 27;
int motor2pin2 = 14;


// === Cấu hình động cơ và bánh xe ===
const float motor_rpm = 208;     // vòng/phút tại 5V (thay đổi theo VDC)
const float wheel_diameter_mm = 65.0;  // đường kính bánh xe (mm)
const float robot_width_cm = 12.0;  // khoảng cách giữa 2 bánh xe (cm)


void setup() {
  // Khởi tạo kết nối Serial
  Serial.begin(115200);

  pinMode(motor1pin1, OUTPUT);
  pinMode(motor1pin2, OUTPUT);
  pinMode(motor2pin1, OUTPUT);
  pinMode(motor2pin2, OUTPUT);

  // Thêm các node
  allNodes.push_back({0, 0.0, 0.0, "STATION", {}});
  allNodes.push_back({1, 50.0, 0.0, "INTERSECTION", {}});
  allNodes.push_back({2, 50.0, 30.0, "ROOM", {}});
  allNodes.push_back({3, 80.0, 30.0, "STATION", {}});

  // Gán neighbor (đồ thị vô hướng)
  allNodes[0].neighbors.push_back({1, 50.0f});
  allNodes[1].neighbors.push_back({0, 50.0f});
  allNodes[1].neighbors.push_back({2, 30.0f});
  allNodes[2].neighbors.push_back({1, 30.0f});
  allNodes[2].neighbors.push_back({3, 30.0f});
  allNodes[3].neighbors.push_back({2, 30.0f});

    // Giả sử ban đầu robot quay sang phải và đang ở tọa độ (0;0)
  currentDirection = EAST;
  currentLocation = 0;
}

void loop() {
  vector<int> path = aStar(currentLocation, 2, allNodes);

  Serial.println("Path:");
  for (int i = 0; i < path.size(); ++i) {
    Serial.print(path[i]);
    if (i < path.size() - 1) Serial.print(" -> ");
  }
  Serial.println();

  followPath(path);
  delay(5000);

  int nearestStationId = findNearestStation(currentLocation, allNodes);
  if (nearestStationId != -1) {
    vector<int> returnPath = aStar(currentLocation, nearestStationId, allNodes);
    Serial.println("Path:");
    for (int i = 0; i < returnPath.size(); ++i) {
      Serial.print(returnPath[i]);
      if (i < returnPath.size() - 1) Serial.print(" -> ");
    }
    Serial.println();
    followPath(returnPath);
    currentLocation = returnPath.back(); // Cập nhật lại vị trí sau khi quay về
  }

  // moveBackward();
  // delay(5000);
  // stopMotors();

  delay(10000);
}


// ====== CÁC HÀM DI CHUYỂN CƠ BẢN ======

void moveForward() {
  Serial.println("Move forward");
  digitalWrite(motor1pin1, LOW);
  digitalWrite(motor1pin2, HIGH);
  digitalWrite(motor2pin1, LOW);
  digitalWrite(motor2pin2, HIGH);
}

void stopMotors() { 
  digitalWrite(motor1pin1, LOW);
  digitalWrite(motor1pin2, LOW);
  digitalWrite(motor2pin1, LOW);
  digitalWrite(motor2pin2, LOW);
}

void turnLeft90() {
  Serial.println("Turn left");
  digitalWrite(motor1pin1, HIGH);
  digitalWrite(motor1pin2, LOW);
  digitalWrite(motor2pin1, LOW);
  digitalWrite(motor2pin2, HIGH);
  delay(450); // thời gian để rẽ tới góc 90 độ
  stopMotors();
}

void turnRight90() {
  Serial.println("Turn right");
  digitalWrite(motor1pin1, LOW);
  digitalWrite(motor1pin2, HIGH);
  digitalWrite(motor2pin1, HIGH);
  digitalWrite(motor2pin2, LOW);
  delay(450); // thời gian để rẽ tới góc 90 độ
  stopMotors();
}

void turnBack() {
  Serial.println("Turn back");
  digitalWrite(motor1pin1, LOW);
  digitalWrite(motor1pin2, HIGH);
  digitalWrite(motor2pin1, HIGH);
  digitalWrite(motor2pin2, LOW);
  delay(900); // thời gian để rẽ tới góc 90 độ
  stopMotors();
}


// ====== HÀM CHÍNH: ĐI THEO CM ======
float calculateDurationMs(float distance_cm, float rpm, float wheel_diameter_mm) {
  // Tính chu vi bánh xe
  float wheel_circumference_cm = 3.1416 * (wheel_diameter_mm / 10.0);  // mm → cm

  // Tính số vòng mỗi giây
  float rps = motor_rpm / 60.0;

  // Tính tốc độ robot: cm/giây
  float speed_cm_per_sec = wheel_circumference_cm * rps;

  // Thời gian cần chạy (ms)
  return (distance_cm / speed_cm_per_sec) * 1000.0;  // ms
}

void moveDistance(float distance_cm) {
  float duration_ms = calculateDurationMs(distance_cm, motor_rpm, wheel_diameter_mm);
  
  moveForward();
  delay(duration_ms);
  stopMotors();
}


// ====== CÁC HÀM ĐIỀU KHIỂN ROBOT ĐI THEO LỘ TRÌNH ======

Direction getDirectionBetween(Node a, Node b) {
  float dx = b.x - a.x;
  float dy = b.y - a.y;

  if (abs(dx) > abs(dy)) {
    return dx > 0 ? EAST : WEST;
  } else {
    return dy > 0 ? NORTH : SOUTH;
  }
}

void rotateToDirection(Direction targetDir) {
  int diff = (targetDir - currentDirection + 4) % 4;

  if (diff == 1) {
    turnLeft90();
  } else if (diff == 3) {
    turnRight90();
  } else if (diff == 2) {
    // turnLeft90();
    // // delay(500); // thêm delay nếu cần
    // turnLeft90();
    turnBack();
  }
  currentDirection = targetDir;
}

void followPath(vector<int> path) {
  for (int i = 0; i < path.size() - 1; ++i) {
    Node& current = allNodes[path[i]];
    Node& next = allNodes[path[i + 1]];

    Direction nextDir = getDirectionBetween(current, next);
    rotateToDirection(nextDir);
    delay(500);

    // Tính khoảng cách cần đi
    float dx = next.x - current.x;
    float dy = next.y - current.y;
    float distance = sqrt(dx * dx + dy * dy);

    moveDistance(distance);
    delay(500); // nghỉ chút giữa các bước đi
  }
  currentLocation = path.back();
}