#include "LoopCharacteristics.h"
#include "Dom.h"
#include "Loop.h"
#include "Liveness.h"
#include "AliasAnalysis.h"
#include "SideEffectAnalysis.h"
#include <iostream>
#include <cmath>

// 使用全局调试开关
extern int DEBUG;

namespace sysy {

// 定义 Pass 的唯一 ID
void *LoopCharacteristicsPass::ID = (void *)&LoopCharacteristicsPass::ID;

void LoopCharacteristicsResult::print() const {
  if (!DEBUG) return;

  std::cout << "\n--- Loop Characteristics Analysis Results for Function: " 
            << AssociatedFunction->getName() << " ---" << std::endl;

  if (CharacteristicsMap.empty()) {
    std::cout << "  No loop characteristics found." << std::endl;
    return;
  }

  // 打印统计信息
  auto stats = getOptimizationStats();
  std::cout << "\n=== Basic Loop Characteristics Statistics ===" << std::endl;
  std::cout << "Total Loops: " << stats.totalLoops << std::endl;
  std::cout << "Counting Loops: " << stats.countingLoops << std::endl;
  std::cout << "Unrolling Candidates: " << stats.unrollingCandidates << std::endl;
  std::cout << "Pure Loops: " << stats.pureLoops << std::endl;
  std::cout << "Local Memory Only Loops: " << stats.localMemoryOnlyLoops << std::endl;
  std::cout << "No Alias Conflict Loops: " << stats.noAliasConflictLoops << std::endl;
  std::cout << "Avg Instructions per Loop: " << stats.avgInstructionCount << std::endl;
  std::cout << "Avg Compute/Memory Ratio: " << stats.avgComputeMemoryRatio << std::endl;

  // 按热度排序并打印循环特征
  auto loopsByHotness = getLoopsByHotness();
  std::cout << "\n=== Loop Characteristics (by hotness) ===" << std::endl;
  
  for (auto* loop : loopsByHotness) {
    auto* chars = getCharacteristics(loop);
    if (!chars) continue;
    
    std::cout << "\n--- Loop: " << loop->getName() << " (Hotness: " 
              << loop->getLoopHotness() << ") ---" << std::endl;
    std::cout << "  Level: " << loop->getLoopLevel() << std::endl;
    std::cout << "  Blocks: " << loop->getLoopSize() << std::endl;
    std::cout << "  Instructions: " << chars->instructionCount << std::endl;
    std::cout << "  Memory Operations: " << chars->memoryOperationCount << std::endl;
    std::cout << "  Compute/Memory Ratio: " << chars->computeToMemoryRatio << std::endl;
    
    // 循环形式
    std::cout << "  Form: ";
    if (chars->isCountingLoop) std::cout << "Counting ";
    if (chars->isSimpleForLoop) std::cout << "SimpleFor ";
    if (chars->isInnermost) std::cout << "Innermost ";
    if (chars->hasComplexControlFlow) std::cout << "Complex ";
    if (chars->isPure) std::cout << "Pure ";
    if (chars->accessesOnlyLocalMemory) std::cout << "LocalMemOnly ";
    if (chars->hasNoMemoryAliasConflicts) std::cout << "NoAliasConflicts ";
    std::cout << std::endl;
    
    // 边界信息
    if (chars->staticTripCount.has_value()) {
      std::cout << "  Static Trip Count: " << *chars->staticTripCount << std::endl;
    }
    if (chars->hasKnownBounds) {
      std::cout << "  Has Known Bounds: Yes" << std::endl;
    }
    
    // 优化机会
    std::cout << "  Optimization Opportunities: ";
    if (chars->benefitsFromUnrolling) 
      std::cout << "Unroll(factor=" << chars->suggestedUnrollFactor << ") ";
    std::cout << std::endl;
    
    // 归纳变量
    if (!chars->basicInductionVars.empty()) {
      std::cout << "  Basic Induction Vars: " << chars->basicInductionVars.size() << std::endl;
    }
    
    // 循环不变量
    if (!chars->loopInvariants.empty()) {
      std::cout << "  Loop Invariants: " << chars->loopInvariants.size() << std::endl;
    }
    if (!chars->invariantInsts.empty()) {
      std::cout << "  Hoistable Instructions: " << chars->invariantInsts.size() << std::endl;
    }
  }
  
  std::cout << "-----------------------------------------------" << std::endl;
}

bool LoopCharacteristicsPass::runOnFunction(Function *F, AnalysisManager &AM) {
  if (F->getBasicBlocks().empty()) {
    CurrentResult = std::make_unique<LoopCharacteristicsResult>(F);
    return false; // 空函数
  }

  if (DEBUG)
    std::cout << "Running LoopCharacteristicsPass on function: " << F->getName() << std::endl;

  // 获取循环分析结果
  auto* loopAnalysisResult = AM.getAnalysisResult<LoopAnalysisResult, LoopAnalysisPass>(F);
  if (!loopAnalysisResult) {
    std::cerr << "Error: LoopAnalysisResult not available for function " << F->getName() << std::endl;
    CurrentResult = std::make_unique<LoopCharacteristicsResult>(F);
    return false;
  }

  // 如果没有循环，直接返回
  if (!loopAnalysisResult->hasLoops()) {
    CurrentResult = std::make_unique<LoopCharacteristicsResult>(F);
    return false;
  }

  // 获取别名分析和副作用分析结果
  auto* aliasAnalysis = AM.getAnalysisResult<AliasAnalysisResult, SysYAliasAnalysisPass>(F);
  auto* sideEffectAnalysis = AM.getAnalysisResult<SideEffectAnalysisResult, SysYSideEffectAnalysisPass>();

  if (DEBUG) {
    if (aliasAnalysis) std::cout << "LoopCharacteristics: Using alias analysis results" << std::endl;
    if (sideEffectAnalysis) std::cout << "LoopCharacteristics: Using side effect analysis results" << std::endl;
  }

  CurrentResult = std::make_unique<LoopCharacteristicsResult>(F);

  // 分析每个循环的特征
  for (const auto& loop_ptr : loopAnalysisResult->getAllLoops()) {
    Loop* loop = loop_ptr.get();
    auto characteristics = std::make_unique<LoopCharacteristics>(loop);
    
    // 执行各种特征分析，传递分析结果
    analyzeLoop(loop, characteristics.get(), AM, aliasAnalysis, sideEffectAnalysis);
    
    // 添加到结果中
    CurrentResult->addLoopCharacteristics(std::move(characteristics));
  }

  if (DEBUG) {
    std::cout << "LoopCharacteristicsPass completed for function: " << F->getName() << std::endl;
    auto stats = CurrentResult->getOptimizationStats();
    std::cout << "Analyzed " << stats.totalLoops << " loops, found " 
              << stats.countingLoops << " counting loops, " 
              << stats.unrollingCandidates << " unroll candidates" << std::endl;
  }

  return false; // 特征分析不修改IR
}

void LoopCharacteristicsPass::analyzeLoop(Loop* loop, LoopCharacteristics* characteristics, 
                                         AnalysisManager &AM, AliasAnalysisResult* aliasAnalysis,
                                         SideEffectAnalysisResult* sideEffectAnalysis) {
  if (DEBUG)
    std::cout << "  Analyzing basic characteristics of loop: " << loop->getName() << std::endl;

  // 按顺序执行基础分析
  computePerformanceMetrics(loop, characteristics);
  analyzeLoopForm(loop, characteristics);
  analyzePurityAndSideEffects(loop, characteristics, sideEffectAnalysis);
  identifyBasicInductionVariables(loop, characteristics);
  identifyBasicLoopInvariants(loop, characteristics);
  analyzeBasicLoopBounds(loop, characteristics);
  analyzeBasicMemoryAccessPatterns(loop, characteristics, aliasAnalysis);
  evaluateBasicOptimizationOpportunities(loop, characteristics);
}

void LoopCharacteristicsPass::computePerformanceMetrics(Loop* loop, LoopCharacteristics* characteristics) {
  size_t totalInsts = 0;
  size_t memoryOps = 0;
  size_t arithmeticOps = 0;

  // 遍历循环中的所有指令
  for (BasicBlock* bb : loop->getBlocks()) {
    for (auto& inst : bb->getInstructions()) {
      totalInsts++;
      
      // 分类指令类型
      if (dynamic_cast<LoadInst*>(inst.get()) || dynamic_cast<StoreInst*>(inst.get())) {
        memoryOps++;
      } else if (dynamic_cast<BinaryInst*>(inst.get())) {
        // 检查是否为算术运算
        auto* binInst = dynamic_cast<BinaryInst*>(inst.get());
        // 简化：假设所有二元运算都是算术运算
        arithmeticOps++;
      }
    }
  }

  characteristics->instructionCount = totalInsts;
  characteristics->memoryOperationCount = memoryOps;
  characteristics->arithmeticOperationCount = arithmeticOps;
  
  // 计算计算与内存操作比率
  if (memoryOps > 0) {
    characteristics->computeToMemoryRatio = static_cast<double>(arithmeticOps) / memoryOps;
  } else {
    characteristics->computeToMemoryRatio = arithmeticOps; // 纯计算循环
  }
}

void LoopCharacteristicsPass::analyzeLoopForm(Loop* loop, LoopCharacteristics* characteristics) {
  // 基本形式判断
  characteristics->isInnermost = loop->isInnermost();
  
  // 检查是否为简单循环 (只有一个回边)
  bool isSimple = loop->isSimpleLoop();
  characteristics->isSimpleForLoop = isSimple;
  
  // 检查复杂控制流 (多个出口表示可能有break/continue)
  auto exitingBlocks = loop->getExitingBlocks();
  characteristics->hasComplexControlFlow = exitingBlocks.size() > 1;
  
  // 初步判断是否为计数循环 (需要更复杂的分析)
  // 简化版本：如果是简单循环且是最内层，很可能是计数循环
  characteristics->isCountingLoop = isSimple && loop->isInnermost() && exitingBlocks.size() == 1;
}

void LoopCharacteristicsPass::analyzePurityAndSideEffects(Loop* loop, LoopCharacteristics* characteristics,
                                                         SideEffectAnalysisResult* sideEffectAnalysis) {
  if (!sideEffectAnalysis) {
    // 没有副作用分析结果，保守处理
    characteristics->isPure = false;
    return;
  }
  
  // 检查循环是否有副作用
  characteristics->isPure = !loop->mayHaveSideEffects(sideEffectAnalysis);
  
  if (DEBUG && characteristics->isPure) {
    std::cout << "    Loop " << loop->getName() << " is identified as PURE (no side effects)" << std::endl;
  }
}

void LoopCharacteristicsPass::analyzeBasicMemoryAccessPatterns(Loop* loop, LoopCharacteristics* characteristics,
                                                              AliasAnalysisResult* aliasAnalysis) {
  if (!aliasAnalysis) {
    // 没有别名分析结果，保守处理
    characteristics->accessesOnlyLocalMemory = false;
    characteristics->hasNoMemoryAliasConflicts = false;
    return;
  }
  
  // 检查是否只访问局部内存
  characteristics->accessesOnlyLocalMemory = !loop->accessesGlobalMemory(aliasAnalysis);
  
  // 检查是否有内存别名冲突
  characteristics->hasNoMemoryAliasConflicts = !loop->hasMemoryAliasConflicts(aliasAnalysis);
  
  if (DEBUG) {
    if (characteristics->accessesOnlyLocalMemory) {
      std::cout << "    Loop " << loop->getName() << " accesses ONLY LOCAL MEMORY" << std::endl;
    }
    if (characteristics->hasNoMemoryAliasConflicts) {
      std::cout << "    Loop " << loop->getName() << " has NO MEMORY ALIAS CONFLICTS" << std::endl;
    }
  }
  
  // 分析基础的内存访问模式
  for (BasicBlock* bb : loop->getBlocks()) {
    for (auto& inst : bb->getInstructions()) {
      if (auto* loadInst = dynamic_cast<LoadInst*>(inst.get())) {
        Value* ptr = loadInst->getPointer();
        
        auto& pattern = characteristics->memoryPatterns[ptr];
        pattern.loadInsts.push_back(loadInst);
        pattern.isArrayParameter = aliasAnalysis->isFunctionParameter(ptr);
        pattern.isGlobalArray = aliasAnalysis->isGlobalArray(ptr);
        pattern.hasConstantIndices = aliasAnalysis->hasConstantAccess(ptr);
        
      } else if (auto* storeInst = dynamic_cast<StoreInst*>(inst.get())) {
        Value* ptr = storeInst->getPointer();
        
        auto& pattern = characteristics->memoryPatterns[ptr];
        pattern.storeInsts.push_back(storeInst);
        pattern.isArrayParameter = aliasAnalysis->isFunctionParameter(ptr);
        pattern.isGlobalArray = aliasAnalysis->isGlobalArray(ptr);
        pattern.hasConstantIndices = aliasAnalysis->hasConstantAccess(ptr);
      }
    }
  }
}

void LoopCharacteristicsPass::identifyBasicInductionVariables(Loop* loop, LoopCharacteristics* characteristics) {
  // 寻找基本归纳变量（简化版本）
  BasicBlock* header = loop->getHeader();
  
  // 遍历循环头的phi指令，寻找基本归纳变量模式
  for (auto& inst : header->getInstructions()) {
    auto* phiInst = dynamic_cast<PhiInst*>(inst.get());
    if (!phiInst) continue;
    
    // 检查phi指令是否符合基本归纳变量模式
    if (isBasicInductionVariable(phiInst, loop)) {
      characteristics->basicInductionVars.push_back(phiInst);
      characteristics->inductionSteps[phiInst] = 1; // 简化：默认步长为1
      
      if (DEBUG)
        std::cout << "    Found basic induction variable: " << phiInst->getName() << std::endl;
    }
  }
}

void LoopCharacteristicsPass::identifyBasicLoopInvariants(Loop* loop, LoopCharacteristics* characteristics) {
  // 收集基础循环不变量（简化版本）
  for (BasicBlock* bb : loop->getBlocks()) {
    for (auto& inst : bb->getInstructions()) {
      Value* val = inst.get();
      
      // 跳过phi指令和终结指令
      if (dynamic_cast<PhiInst*>(val)) continue;
      if (auto* instPtr = dynamic_cast<Instruction*>(val)) {
        if (instPtr->isTerminator()) continue;
      }
      
      if (isBasicLoopInvariant(val, loop)) {
        characteristics->loopInvariants.insert(val);
        characteristics->invariantInsts.insert(static_cast<Instruction*>(val));
        
        if (DEBUG)
          std::cout << "    Found basic loop invariant: " << val->getName() << std::endl;
      }
    }
  }
}

void LoopCharacteristicsPass::analyzeBasicLoopBounds(Loop* loop, LoopCharacteristics* characteristics) {
  // 简化的基础边界分析
  // 检查是否有静态可确定的循环次数（简化版本）
  if (characteristics->isCountingLoop && !characteristics->basicInductionVars.empty()) {
    // 简化：如果是计数循环且有基本归纳变量，尝试确定循环次数
    if (characteristics->instructionCount < 10) {
      characteristics->staticTripCount = 100; // 简化估计
      characteristics->hasKnownBounds = true;
      
      if (DEBUG) {
        std::cout << "    Estimated static trip count: " << *characteristics->staticTripCount << std::endl;
      }
    }
  }
}

void LoopCharacteristicsPass::evaluateBasicOptimizationOpportunities(Loop* loop, LoopCharacteristics* characteristics) {
  // 评估基础循环展开机会
  characteristics->benefitsFromUnrolling = 
    characteristics->isInnermost && 
    characteristics->instructionCount > 3 && 
    characteristics->instructionCount < 50 &&
    !characteristics->hasComplexControlFlow;
    
  if (characteristics->benefitsFromUnrolling) {
    // 基于循环体大小估算展开因子
    if (characteristics->instructionCount <= 5) characteristics->suggestedUnrollFactor = 8;
    else if (characteristics->instructionCount <= 10) characteristics->suggestedUnrollFactor = 4;
    else if (characteristics->instructionCount <= 20) characteristics->suggestedUnrollFactor = 2;
    else characteristics->suggestedUnrollFactor = 1;
  }
    
  if (DEBUG) {
    if (characteristics->benefitsFromUnrolling) {
      std::cout << "    Loop " << loop->getName() << " benefits from UNROLLING (factor=" 
                << characteristics->suggestedUnrollFactor << ")" << std::endl;
    }
  }
}

// ========== 辅助方法实现 ==========

bool LoopCharacteristicsPass::isBasicInductionVariable(Value* val, Loop* loop) {
  // 简化的基础归纳变量检测
  auto* phiInst = dynamic_cast<PhiInst*>(val);
  if (!phiInst) return false;
  
  // 检查phi指令是否在循环头
  if (phiInst->getParent() != loop->getHeader()) return false;
  
  // 检查是否有来自循环内的更新
  for (auto& [incomingBB, incomingVal] : phiInst->getIncomingValues()) {
    if (loop->contains(incomingBB)) {
      return true; // 简化：有来自循环内的值就认为是基础归纳变量
    }
  }
  
  return false;
}

bool LoopCharacteristicsPass::isBasicLoopInvariant(Value* val, Loop* loop) {
  auto* inst = dynamic_cast<Instruction*>(val);
  if (!inst) return true; // 非指令（如常量）认为是不变的
  
  // 如果指令不在循环内定义，则是不变的
  if (!loop->contains(inst->getParent())) {
    return true;
  }
  
  // 简化的基础不变量检测：load指令且指针是循环外的
  if (auto* loadInst = dynamic_cast<LoadInst*>(inst)) {
    Value* ptr = loadInst->getPointer();
    return isBasicLoopInvariant(ptr, loop);
  }
  
  // 保守：对于其他指令，认为是变化的
  return false;
}

bool LoopCharacteristicsPass::hasSimpleMemoryPattern(Loop* loop) {
  // 检查是否有简单的内存访问模式
  return true; // 暂时简化处理
}

} // namespace sysy