#include "UWBHelper.h"
#include <math.h>

UWBHelper::UWBHelper(HardwareSerial* serial, int reset) {
    uwbSerial = serial;
    resetPin = reset;
}

// ===== BASIC COMMANDS (3.2-3.6) =====

bool UWBHelper::testConnection() {
    String response = sendCommand("AT?", 2000);
    return isResponseOK(response);
}

String UWBHelper::getVersion() {
    return sendCommand("AT+GETVER?", 2000);
}

bool UWBHelper::restart() {
    String response = sendCommand("AT+RESTART", 2000);
    return isResponseOK(response);
}

bool UWBHelper::restore() {
    String response = sendCommand("AT+RESTORE", 2000);
    return isResponseOK(response);
}

bool UWBHelper::save() {
    String response = sendCommand("AT+SAVE", 2000);
    return isResponseOK(response);
}

// ===== CONFIGURATION COMMANDS (3.7-3.8) =====

bool UWBHelper::setConfig(int deviceId, int role, int dataRate, int rangeFilter) {
    String cmd = "AT+SETCFG=" + String(deviceId) + "," + 
                 String(role) + "," + 
                 String(dataRate) + "," + 
                 String(rangeFilter);
    
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getConfig() {
    return sendCommand("AT+GETCFG?", 2000);
}

// ===== ANTENNA COMMANDS (3.9-3.10) =====

bool UWBHelper::setAntennaDelay(int delay) {
    String cmd = "AT+SETANT=" + String(delay);
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getAntennaDelay() {
    return sendCommand("AT+GETANT?", 2000);
}

// ===== CAPACITY COMMANDS (3.11-3.12) =====

bool UWBHelper::setCapacity(int tagCount, int timeSlot, int extMode) {
    String cmd = "AT+SETCAP=" + String(tagCount) + "," + 
                 String(timeSlot) + "," + 
                 String(extMode);
    
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getCapacity() {
    return sendCommand("AT+GETCAP?", 2000);
}

// ===== REPORTING COMMANDS (3.13-3.14) =====

bool UWBHelper::setReporting(bool enable) {
    String cmd = "AT+SETRPT=" + String(enable ? 1 : 0);
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getReporting() {
    return sendCommand("AT+GETRPT?", 2000);
}

// ===== SLEEP COMMAND (3.16) =====

bool UWBHelper::setSleep(int sleepTime) {
    String cmd = "AT+SLEEP=" + String(sleepTime);
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

// ===== POWER COMMANDS (3.17-3.18) =====

bool UWBHelper::setPower(String powerValue) {
    String cmd = "AT+SETPOW=" + powerValue;
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getPower() {
    return sendCommand("AT+GETPOW?", 2000);
}

// ===== DATA COMMANDS (3.19-3.20) =====

bool UWBHelper::sendData(int length, String data) {
    String cmd = "AT+DATA=" + String(length) + "," + data;
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::receiveData() {
    return sendCommand("AT+RDATA", 2000);
}

// ===== NETWORK COMMANDS (3.21-3.22) =====

bool UWBHelper::setNetwork(int networkId) {
    String cmd = "AT+SETPAN=" + String(networkId);
    String response = sendCommand(cmd, 2000);
    return isResponseOK(response);
}

String UWBHelper::getNetwork() {
    return sendCommand("AT+GETPAN?", 2000);
}

// ===== LEGACY WRAPPER FUNCTIONS =====

bool UWBHelper::begin() {
    pinMode(resetPin, OUTPUT);
    digitalWrite(resetPin, HIGH);
    
    uwbSerial->begin(115200);
    delay(1000);
    
    // Test communication
    return testConnection();
}

void UWBHelper::reset() {
    digitalWrite(resetPin, LOW);
    delay(100);
    digitalWrite(resetPin, HIGH);
    delay(1000);
}

bool UWBHelper::configureDevice(int deviceId, bool isAnchor, int dataRate, int rangeFilter) {
    return setConfig(deviceId, isAnchor ? 1 : 0, dataRate, rangeFilter);
}

bool UWBHelper::enableReporting(bool enable) {
    return setReporting(enable);
}

bool UWBHelper::saveConfiguration() {
    return save();
}

bool UWBHelper::restartDevice() {
    return restart();
}

// ===== COMMUNICATION =====

String UWBHelper::sendCommand(String command, int timeout) {
    String response = "";
    
    Serial.println("Sending: " + command);
    uwbSerial->println(command);
    
    unsigned long startTime = millis();
    while ((millis() - startTime) < timeout) {
        while (uwbSerial->available()) {
            char c = uwbSerial->read();
            if (c == '\n' || c == '\r') {
                if (response.length() > 0) {
                    Serial.println("Response: " + response);
                    return response;
                }
            } else {
                response += c;
            }
        }
        delay(1);
    }
    
    if (response.length() > 0) {
        Serial.println("Response: " + response);
    }
    return response;
}

// ===== RANGE DATA PARSING =====

bool UWBHelper::parseRangeData(String data, DeviceData devices[], int maxDevices, bool isAnchor) {
    if (!data.startsWith("AT+RANGE=")) {
        return false;
    }
    
    // Parse tid (Tag ID) - this is who's ranging, not who we're ranging to
    int tidIndex = data.indexOf("tid:");
    if (tidIndex == -1) return false;
    
    // Parse range data - array of distances to each anchor
    int rangeIndex = data.indexOf("range:(");
    if (rangeIndex == -1) return false;
    
    int rangeStart = rangeIndex + 7;
    int rangeEnd = data.indexOf(')', rangeStart);
    if (rangeEnd == -1) return false;
    String rangeData = data.substring(rangeStart, rangeEnd);
    
    // Parse RSSI data - array of RSSI values from each anchor
    int rssiIndex = data.indexOf("rssi:(");
    if (rssiIndex == -1) return false;
    
    int rssiStart = rssiIndex + 6;
    int rssiEnd = data.indexOf(')', rssiStart);
    if (rssiEnd == -1) return false;
    String rssiData = data.substring(rssiStart, rssiEnd);
    
    // Parse anchor IDs - which anchors are active
    int ancidIndex = data.indexOf("ancid:(");
    if (ancidIndex == -1) return false;
    
    int ancidStart = ancidIndex + 7;
    int ancidEnd = data.indexOf(')', ancidStart);
    if (ancidEnd == -1) return false;
    String ancidData = data.substring(ancidStart, ancidEnd);
    
    // Parse arrays and match anchor IDs with their distances and RSSI
    float ranges[8];
    float rssiValues[8];
    int anchorIds[8];
    
    // Parse range values
    int startIdx = 0;
    for (int i = 0; i < 8; i++) {
        int commaIdx = rangeData.indexOf(',', startIdx);
        if (commaIdx == -1) commaIdx = rangeData.length();
        ranges[i] = rangeData.substring(startIdx, commaIdx).toFloat() / 100.0; // Convert cm to m
        startIdx = commaIdx + 1;
        if (startIdx >= rangeData.length()) break;
    }
    
    // Parse RSSI values
    startIdx = 0;
    for (int i = 0; i < 8; i++) {
        int commaIdx = rssiData.indexOf(',', startIdx);
        if (commaIdx == -1) commaIdx = rssiData.length();
        rssiValues[i] = rssiData.substring(startIdx, commaIdx).toFloat();
        startIdx = commaIdx + 1;
        if (startIdx >= rssiData.length()) break;
    }
    
    // Parse anchor IDs
    startIdx = 0;
    for (int i = 0; i < 8; i++) {
        int commaIdx = ancidData.indexOf(',', startIdx);
        if (commaIdx == -1) commaIdx = ancidData.length();
        anchorIds[i] = ancidData.substring(startIdx, commaIdx).toInt();
        startIdx = commaIdx + 1;
        if (startIdx >= ancidData.length()) break;
    }
    
    // Parse tid (Tag ID) to know which device is reporting ranges
    int commaPos = data.indexOf(',', tidIndex);
    if (commaPos == -1) return false;
    int reportingDeviceId = data.substring(tidIndex + 4, commaPos).toInt();
    
    if (isAnchor) {
        // ANCHOR BEHAVIOR: When tag reports (tid >= 1), only store the tag itself
        if (reportingDeviceId >= 1) {
            // Find first non-zero distance and RSSI for this reporting tag
            float tagDistance = 0.0;
            float tagRssi = 0.0;
            
            for (int i = 0; i < 8; i++) {
                if (ranges[i] > 0) {
                    tagDistance = ranges[i];
                    tagRssi = rssiValues[i];
                    break;
                }
            }
            
            if (tagDistance > 0) {
                // Find or create slot for this tag
                int deviceSlot = -1;
                for (int j = 0; j < maxDevices; j++) {
                    if (devices[j].deviceId == reportingDeviceId || !devices[j].active) {
                        deviceSlot = j;
                        break;
                    }
                }
                
                if (deviceSlot >= 0) {
                    devices[deviceSlot].deviceId = reportingDeviceId;
                    devices[deviceSlot].distance = tagDistance;
                    devices[deviceSlot].rssi = tagRssi;
                    devices[deviceSlot].lastUpdate = millis();
                    devices[deviceSlot].active = true;
                    return true;
                }
            }
        }
    } else {
        // TAG BEHAVIOR: When tag reports (tid >= 1), store all anchors in the range data
        if (reportingDeviceId >= 1) {
            bool dataUpdated = false;
            for (int i = 0; i < 8; i++) {
                if (anchorIds[i] >= 0 && ranges[i] > 0) {
                    // Find or create slot for this anchor
                    int deviceSlot = -1;
                    for (int j = 0; j < maxDevices; j++) {
                        if (devices[j].deviceId == anchorIds[i] || !devices[j].active) {
                            deviceSlot = j;
                            break;
                        }
                    }
                    
                    if (deviceSlot >= 0) {
                        devices[deviceSlot].deviceId = anchorIds[i];
                        devices[deviceSlot].distance = ranges[i];
                        devices[deviceSlot].rssi = rssiValues[i];
                        devices[deviceSlot].lastUpdate = millis();
                        devices[deviceSlot].active = true;
                        dataUpdated = true;
                    }
                }
            }
            return dataUpdated;
        }
    }
    
    return false;
}

bool UWBHelper::parseDetailedRangeData(String data, RangeResult* result) {
    if (!data.startsWith("AT+RANGE=")) {
        return false;
    }
    
    // Parse tid
    int tidIndex = data.indexOf("tid:");
    if (tidIndex == -1) return false;
    int commaPos = data.indexOf(',', tidIndex);
    if (commaPos == -1) return false;
    result->tagId = data.substring(tidIndex + 4, commaPos).toInt();
    
    // Parse timer (if present)
    int timerIndex = data.indexOf("timer:");
    if (timerIndex != -1) {
        int timerComma = data.indexOf(',', timerIndex);
        if (timerComma == -1) timerComma = data.indexOf(',', timerIndex);
        result->timer = data.substring(timerIndex + 6, timerComma).toInt();
    }
    
    // Parse timerSys (if present)
    int timerSysIndex = data.indexOf("timerSys:");
    if (timerSysIndex != -1) {
        int timerSysComma = data.indexOf(',', timerSysIndex);
        if (timerSysComma == -1) timerSysComma = data.indexOf(',', timerSysIndex);
        result->timerSys = data.substring(timerSysIndex + 9, timerSysComma).toInt();
    }
    
    // Parse mask
    int maskIndex = data.indexOf("mask:");
    if (maskIndex != -1) {
        int maskComma = data.indexOf(',', maskIndex);
        result->mask = data.substring(maskIndex + 5, maskComma).toInt();
    }
    
    // Parse sequence
    int seqIndex = data.indexOf("seq:");
    if (seqIndex != -1) {
        int seqComma = data.indexOf(',', seqIndex);
        result->sequence = data.substring(seqIndex + 4, seqComma).toInt();
    }
    
    // Parse ranges
    int rangeIndex = data.indexOf("range:(");
    if (rangeIndex != -1) {
        int rangeStart = rangeIndex + 7;
        int rangeEnd = data.indexOf(')', rangeStart);
        String rangeData = data.substring(rangeStart, rangeEnd);
        
        int startIdx = 0;
        for (int i = 0; i < 8; i++) {
            int commaIdx = rangeData.indexOf(',', startIdx);
            if (commaIdx == -1) commaIdx = rangeData.length();
            
            result->ranges[i] = rangeData.substring(startIdx, commaIdx).toFloat() / 100.0;
            startIdx = commaIdx + 1;
            if (startIdx >= rangeData.length()) break;
        }
    }
    
    // Parse RSSI
    int rssiIndex = data.indexOf("rssi:(");
    if (rssiIndex != -1) {
        int rssiStart = rssiIndex + 6;
        int rssiEnd = data.indexOf(')', rssiStart);
        String rssiData = data.substring(rssiStart, rssiEnd);
        
        int startIdx = 0;
        for (int i = 0; i < 8; i++) {
            int commaIdx = rssiData.indexOf(',', startIdx);
            if (commaIdx == -1) commaIdx = rssiData.length();
            
            result->rssi[i] = rssiData.substring(startIdx, commaIdx).toFloat();
            startIdx = commaIdx + 1;
            if (startIdx >= rssiData.length()) break;
        }
    }
    
    // Parse anchor IDs
    int ancidIndex = data.indexOf("ancid:(");
    if (ancidIndex != -1) {
        int ancidStart = ancidIndex + 7;
        int ancidEnd = data.indexOf(')', ancidStart);
        String ancidData = data.substring(ancidStart, ancidEnd);
        
        int startIdx = 0;
        for (int i = 0; i < 8; i++) {
            int commaIdx = ancidData.indexOf(',', startIdx);
            if (commaIdx == -1) commaIdx = ancidData.length();
            
            result->anchorIds[i] = ancidData.substring(startIdx, commaIdx).toInt();
            startIdx = commaIdx + 1;
            if (startIdx >= ancidData.length()) break;
        }
    }
    
    return true;
}

// ===== ADVANCED FUNCTIONS =====

bool UWBHelper::requestAnchorPosition(int anchorId, AnchorPosition* position) {
    // Custom command to request anchor position
    String cmd = "AT+GETPOS=" + String(anchorId);
    String response = sendCommand(cmd, 2000);
    
    if (response.indexOf("POS=") >= 0) {
        // Parse response format: POS=id,x,y,confidence
        int idStart = response.indexOf("=") + 1;
        int comma1 = response.indexOf(",", idStart);
        int comma2 = response.indexOf(",", comma1 + 1);
        int comma3 = response.indexOf(",", comma2 + 1);
        
        if (comma1 > 0 && comma2 > 0 && comma3 > 0) {
            position->anchorId = response.substring(idStart, comma1).toInt();
            position->x = response.substring(comma1 + 1, comma2).toFloat();
            position->y = response.substring(comma2 + 1, comma3).toFloat();
            position->confidence = response.substring(comma3 + 1).toFloat();
            position->valid = true;
            return true;
        }
    }
    
    position->valid = false;
    return false;
}

bool UWBHelper::calculatePosition(DeviceData devices[], int deviceCount, float* x, float* y) {
    // Simple trilateration with first 3 active devices
    int activeCount = 0;
    DeviceData activeDevices[3];
    
    for (int i = 0; i < deviceCount && activeCount < 3; i++) {
        if (devices[i].active) {
            activeDevices[activeCount] = devices[i];
            activeCount++;
        }
    }
    
    if (activeCount < 3) return false;
    
    // For now, assume anchors are at fixed positions (would need anchor position data)
    // This is a placeholder - real implementation would use actual anchor coordinates
    return false;
}

// ===== UTILITY FUNCTIONS =====

bool UWBHelper::isResponseOK(String response) {
    return response.indexOf("OK") >= 0;
}

void UWBHelper::printDiagnostics() {
    Serial.println("=== UWB Diagnostics ===");
    Serial.println("Version: " + getVersion());
    Serial.println("Config: " + getConfig());
    Serial.println("Capacity: " + getCapacity());
    Serial.println("Antenna: " + getAntennaDelay());
    Serial.println("Power: " + getPower());
    Serial.println("Network: " + getNetwork());
    Serial.println("Reporting: " + getReporting());
}

// ===== DISTANCE FILTER IMPLEMENTATION =====

DistanceFilter::DistanceFilter() : index(0), filled(false) {
    for (int i = 0; i < FILTER_SIZE; i++) readings[i] = 0;
}

float DistanceFilter::addReading(float distance) {
    readings[index] = distance;
    index = (index + 1) % FILTER_SIZE;
    if (index == 0) filled = true;
    
    return getFilteredValue();
}

float DistanceFilter::getFilteredValue() {
    if (!filled && index == 0) return readings[0];
    
    // Median filter
    float sorted[FILTER_SIZE];
    int count = filled ? FILTER_SIZE : index;
    
    for (int i = 0; i < count; i++) {
        sorted[i] = readings[i];
    }
    
    // Simple bubble sort
    for (int i = 0; i < count - 1; i++) {
        for (int j = 0; j < count - i - 1; j++) {
            if (sorted[j] > sorted[j + 1]) {
                float temp = sorted[j];
                sorted[j] = sorted[j + 1];
                sorted[j + 1] = temp;
            }
        }
    }
    
    return sorted[count / 2]; // Return median
}

void DistanceFilter::reset() {
    index = 0;
    filled = false;
    for (int i = 0; i < FILTER_SIZE; i++) readings[i] = 0;
}

// ===== POSITION CALCULATOR IMPLEMENTATION =====

bool PositionCalculator::trilaterate(float x1, float y1, float r1,
                                   float x2, float y2, float r2,
                                   float x3, float y3, float r3,
                                   float* x, float* y) {
    // Trilateration algorithm
    float A = 2 * (x2 - x1);
    float B = 2 * (y2 - y1);
    float C = pow(r1, 2) - pow(r2, 2) - pow(x1, 2) + pow(x2, 2) - pow(y1, 2) + pow(y2, 2);
    float D = 2 * (x3 - x2);
    float E = 2 * (y3 - y2);
    float F = pow(r2, 2) - pow(r3, 2) - pow(x2, 2) + pow(x3, 2) - pow(y2, 2) + pow(y3, 2);
    
    float denominator = A * E - B * D;
    if (abs(denominator) < 0.001) return false; // Points are collinear
    
    *x = (C * E - F * B) / denominator;
    *y = (A * F - D * C) / denominator;
    
    return true;
}

bool PositionCalculator::multilaterate(AnchorPosition anchors[], float distances[], int count, float* x, float* y) {
    if (count < 3) return false;
    
    // Use least squares method for more than 3 anchors
    if (count == 3) {
        return trilaterate(anchors[0].x, anchors[0].y, distances[0],
                          anchors[1].x, anchors[1].y, distances[1],
                          anchors[2].x, anchors[2].y, distances[2],
                          x, y);
    }
    
    // For more than 3 anchors, use weighted least squares (simplified version)
    float sumX = 0, sumY = 0, sumW = 0;
    
    for (int i = 0; i < count - 1; i++) {
        for (int j = i + 1; j < count; j++) {
            for (int k = j + 1; k < count; k++) {
                float tempX, tempY;
                if (trilaterate(anchors[i].x, anchors[i].y, distances[i],
                               anchors[j].x, anchors[j].y, distances[j],
                               anchors[k].x, anchors[k].y, distances[k],
                               &tempX, &tempY)) {
                    float weight = anchors[i].confidence * anchors[j].confidence * anchors[k].confidence;
                    sumX += tempX * weight;
                    sumY += tempY * weight;
                    sumW += weight;
                }
            }
        }
    }
    
    if (sumW > 0) {
        *x = sumX / sumW;
        *y = sumY / sumW;
        return true;
    }
    
    return false;
}

float PositionCalculator::calculateDistance(float x1, float y1, float x2, float y2) {
    return sqrt(pow(x2 - x1, 2) + pow(y2 - y1, 2));
}