#include <map>
#include <unordered_set>
#include <vector>
#include <queue>
#include <iostream>
#include <algorithm>

#include "../../Utils/errors.h"
#include "../../Utils/SgUtils.h"
#include "../../GraphCall/graph_calls.h"
#include "../../GraphCall/graph_calls_func.h"
#include "../../CFGraph/CFGraph.h"
#include "../../CFGraph/IR.h"
#include "../../GraphLoop/graph_loops.h"
#include "swap_operators.h"

using namespace std;

string getNameByArg(SAPFOR::Argument* arg);
SgSymbol* getSybolByArg(SAPFOR::Argument* arg);

static vector<SAPFOR::IR_Block*> findInstructionsFromOperator(SgStatement* st, vector<SAPFOR::BasicBlock*> Blocks) {
    vector<SAPFOR::IR_Block*> result;
    string filename = st->fileName();
    
    for (auto& block: Blocks) {
        vector<SAPFOR::IR_Block*> instructionsInBlock = block->getInstructions();
        for (auto& instruction: instructionsInBlock) {
            SgStatement* curOperator = instruction->getInstruction()->getOperator();
            // Match by line number to find corresponding IR instruction
            if (curOperator->lineNumber() == st->lineNumber()) {
                result.push_back(instruction);
            }
        }
    }
    return result;
}

unordered_set<int> loop_tags = {FOR_NODE};
unordered_set<int> control_tags = {IF_NODE, ELSEIF_NODE, DO_WHILE_NODE, WHILE_NODE};
unordered_set<int> control_end_tags = {CONTROL_END};

struct OperatorInfo {
    SgStatement* stmt;
    set<string> usedVars;
    set<string> definedVars;
    int lineNumber;
    bool isMovable;
    OperatorInfo(SgStatement* s) : stmt(s), lineNumber(s->lineNumber()), isMovable(true) {}
};

static vector<OperatorInfo> analyzeOperatorsInLoop(SgForStmt* loop, vector<SAPFOR::BasicBlock*> blocks, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR) {
    vector<OperatorInfo> operators;
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    
    SgStatement* current = loopStart;
    while (current && current != loopEnd) {
        if (isSgExecutableStatement(current)) {
            OperatorInfo opInfo(current);
            
            vector<SAPFOR::IR_Block*> irBlocks = findInstructionsFromOperator(current, blocks);
            for (auto irBlock : irBlocks) {
                SAPFOR::Instruction* instr = irBlock->getInstruction();
                
                if (instr->getArg1()) {
                    string varName = getNameByArg(instr->getArg1());
                    if (!varName.empty()) {
                        opInfo.usedVars.insert(varName);
                    }
                }
                if (instr->getArg2()) {
                    string varName = getNameByArg(instr->getArg2());
                    if (!varName.empty()) {
                        opInfo.usedVars.insert(varName);
                    }
                }
                if (instr->getResult()) {
                    string varName = getNameByArg(instr->getResult());
                    if (!varName.empty()) {
                        opInfo.definedVars.insert(varName);
                    }
                }
            }
            
            if (control_tags.find(current->variant()) != control_tags.end()) {
                opInfo.isMovable = false;
            }
            
            operators.push_back(opInfo);
        }
        current = current->lexNext();
    }
    
    return operators;
}

static map<string, vector<SgStatement*>> findVariableDefinitions(SgForStmt* loop, vector<OperatorInfo>& operators) {
    map<string, vector<SgStatement*>> varDefinitions;
    
    for (auto& op : operators) {
        for (const string& var : op.definedVars) {
            varDefinitions[var].push_back(op.stmt);
        }
    }
    
    return varDefinitions;
}

static int calculateDistance(SgStatement* from, SgStatement* to) {
    if (!from || !to) return INT_MAX;
    return abs(to->lineNumber() - from->lineNumber());
}

static SgStatement* findBestPosition(SgStatement* operatorStmt, vector<OperatorInfo>& operators, map<string, vector<SgStatement*>>& varDefinitions) {
    OperatorInfo* opInfo = nullptr;
    for (auto& op : operators) {
        if (op.stmt == operatorStmt) {
            opInfo = &op;
            break;
        }
    }
    
    if (!opInfo || !opInfo->isMovable) return nullptr;
    
    SgStatement* bestPos = nullptr;
    int minDistance = INT_MAX;
    
    for (const string& usedVar : opInfo->usedVars) {
        if (varDefinitions.find(usedVar) != varDefinitions.end()) {
            for (SgStatement* defStmt : varDefinitions[usedVar]) {
                int distance = calculateDistance(operatorStmt, defStmt);
                if (distance < minDistance) {
                    minDistance = distance;
                    bestPos = defStmt;
                }
            }
        }
    }
    
    return bestPos;
}

static bool canMoveTo(SgStatement* from, SgStatement* to, SgForStmt* loop) {
    if (!from || !to || from == to) return false;
    
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    
    if (to->lineNumber() < loopStart->lineNumber() || to->lineNumber() > loopEnd->lineNumber()) {
        return false;
    }
    
    SgStatement* current = from;
    while (current && current != loopEnd) {
        if (control_tags.find(current->variant()) != control_tags.end()) {
            return false;
        }
        if (current == to) break;
        current = current->lexNext();
    }
    
    return true;
}

static vector<SgStatement*> optimizeOperatorOrder(SgForStmt* loop, vector<OperatorInfo>& operators, map<string, vector<SgStatement*>>& varDefinitions) {
    vector<SgStatement*> newOrder;
    vector<bool> moved(operators.size(), false);
    
    for (size_t i = 0; i < operators.size(); i++) {
        if (moved[i] || !operators[i].isMovable) {
            newOrder.push_back(operators[i].stmt);
            moved[i] = true;
            continue;
        }
        
        SgStatement* bestPos = findBestPosition(operators[i].stmt, operators, varDefinitions);
        
        if (bestPos && canMoveTo(operators[i].stmt, bestPos, loop)) {
            bool inserted = false;
            for (size_t j = 0; j < newOrder.size(); j++) {
                if (newOrder[j] == bestPos) {
                    newOrder.insert(newOrder.begin() + j + 1, operators[i].stmt);
                    inserted = true;
                    break;
                }
            }
            if (!inserted) {
                newOrder.push_back(operators[i].stmt);
            }
        } else {
            newOrder.push_back(operators[i].stmt);
        }
        moved[i] = true;
    }
    
    return newOrder;
}

static bool applyOperatorReordering(SgForStmt* loop, vector<SgStatement*>& newOrder) {
    if (!loop || newOrder.empty()) return false;
    
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    
    vector<SgStatement*> extractedStatements;
    vector<char*> savedComments;
    
    for (SgStatement* stmt : newOrder) {
        if (stmt && stmt != loop && stmt != loopEnd) {
            savedComments.push_back(stmt->comments() ? strdup(stmt->comments()) : nullptr);
            SgStatement* extracted = stmt->extractStmt();
            if (extracted) {
                extractedStatements.push_back(extracted);
            }
        }
    }
    
    SgStatement* currentPos = loop;
    int lineCounter = loop->lineNumber() + 1;
    
    for (size_t i = 0; i < extractedStatements.size(); i++) {
        SgStatement* stmt = extractedStatements[i];
        if (stmt) {
            if (i < savedComments.size() && savedComments[i]) {
                stmt->setComments(savedComments[i]);
            }
            stmt->setlineNumber(lineCounter++);
            currentPos->insertStmtAfter(*stmt, *loop);
            currentPos = stmt;
        }
    }
    
    for (char* comment : savedComments) {
        if (comment) {
            free(comment);
        }
    }
    
    if (currentPos && currentPos->lexNext() != loopEnd) {
        currentPos->setLexNext(*loopEnd);
    }
    
    return true;
}

vector<SAPFOR::BasicBlock*> findFuncBlocksByFuncStatement(SgStatement *st, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR) {
    vector<SAPFOR::BasicBlock*> result;
    Statement* forSt = (Statement*)st;
    for (auto& func: FullIR) {
        if (func.first -> funcPointer -> getCurrProcessFile() == forSt -> getCurrProcessFile()
            && func.first -> funcPointer -> lineNumber() == forSt -> lineNumber())
            result = func.second;
    }
    return result;
}

map<SgForStmt*, vector<SAPFOR::BasicBlock*>> findAndAnalyzeLoops(SgStatement *st, vector<SAPFOR::BasicBlock*> blocks) {
    map<SgForStmt*, vector<SAPFOR::BasicBlock*>> result;
    SgStatement *lastNode = st->lastNodeOfStmt();
    while (st && st != lastNode) {
        if (loop_tags.find(st -> variant()) != loop_tags.end()) {
            SgForStmt *forSt = (SgForStmt*)st;
            SgStatement *loopBody = forSt -> body();
            SgStatement *lastLoopNode = st->lastNodeOfStmt();
            unordered_set<int> blocks_nums;
            while (loopBody && loopBody != lastLoopNode) {
                SAPFOR::IR_Block* IR = findInstructionsFromOperator(loopBody, blocks).front();
                if (blocks_nums.find(IR -> getBasicBlock() -> getNumber()) == blocks_nums.end()) {
                    result[forSt].push_back(IR -> getBasicBlock());
                    blocks_nums.insert(IR -> getBasicBlock() -> getNumber());
                }
                loopBody = loopBody -> lexNext();
            }
            std::sort(result[forSt].begin(), result[forSt].end());
        }
        st = st -> lexNext();
    }
    return result;
}

void runSwapOperators(SgFile *file, std::map<std::string, std::vector<LoopGraph*>>& loopGraph, std::map<FuncInfo*, std::vector<SAPFOR::BasicBlock*>>& FullIR, int& countOfTransform) {
    countOfTransform += 1;

    const int funcNum = file -> numberOfFunctions();
    for (int i = 0; i < funcNum; ++i) {
        SgStatement *st = file -> functions(i);
        vector<SAPFOR::BasicBlock*> blocks = findFuncBlocksByFuncStatement(st, FullIR);
        map<SgForStmt*, vector<SAPFOR::BasicBlock*>> loopsMapping = findAndAnalyzeLoops(st, blocks);
        
        for (pair<SgForStmt*, vector<SAPFOR::BasicBlock*>> loopForAnalyze: loopsMapping) {
            vector<OperatorInfo> operators = analyzeOperatorsInLoop(loopForAnalyze.first, loopForAnalyze.second, FullIR);
            map<string, vector<SgStatement*>> varDefinitions = findVariableDefinitions(loopForAnalyze.first, operators);
            
            vector<SgStatement*> newOrder = optimizeOperatorOrder(loopForAnalyze.first, operators, varDefinitions);
            applyOperatorReordering(loopForAnalyze.first, newOrder);
        }
    }
}