from adder import *
from alu_controller import *
from alu import *
from binandint import *
from branch_unit import *
from data_mem import *
from forwarding_unit import *
from hazard_detector import *
from Imm_gen import *
from insn_mem import *
from mux import *
from mux4 import *
from proc_controller import *
from reg_file import reg_file

IM = insn_mem()
PC = 0


def IF(BrPC, Brflush, stall, PC, IM):
    PCplus4 = adder(PC, 4)
    if Brflush == '1':
        nPC = BrPC
    else:
        nPC = PC + 4
    insn = IM.fetch(nPC)
    return nPC, insn


REG = reg_file()


def ID(PC, insn, RegWrite, Brflush, stall, REG, WB_data):
    funct7 = insn[0:7]
    rs2 = BintoUInt(insn[7:12])
    rs1 = BintoUInt(insn[12:17])
    funct3 = insn[17:20]
    rd = BintoUInt(insn[20:25])
    opcode = insn[25:32]
    ALUSrc, MemtoReg, RegWritex, MemRead, MemWrite, ALUOp, Branch, JalrSel, RWSel = proc_controller(
        opcode)
    ImmG = Imm_gen(insn)
    REG.write(RegWrite, rd, WB_data)
    RD1, RD2 = REG.read(rs1, rs2)
    return ALUSrc, MemtoReg, RegWritex, MemRead, MemWrite, ALUOp, Branch, JalrSel, RWSel, PC, RD1, RD2, ImmG, rs1, rs2, rd, funct3, funct7


def EX(ALUSrc, MemtoReg, RegWrite, MemRead, MemWrite, ALUOp, Branch, JalrSel,
       RWSel, PC, RD1, RD2, ImmG, rs1, rs2, rd, funct3, funct7, lastrs1,
       lastrs2, forwardA, forwardB, WB_data, alu_out):
    stall = hazard_detector(lastrs1, lastrs2, rd, MemRead)
    tempA = mux4(RD1, alu_out, WB_data, 0, forwardA)
    tempB = mux4(RD2, alu_out, WB_data, 0, forwardB)
    tempB2 = mux(tempB, ImmG, ALUSrc)
    alu_ctrl = alu_controller(ALUOp, funct7, funct3)
    ALU_result = alu(tempA, tempB2, alu_ctrl)
    pc_plus_imm, pc_plus_4, branch_target, pc_sel = branch_unit(
        PC, ImmG, JalrSel, Branch, ALU_result)
    return MemtoReg, RegWrite, MemWrite, RWSel, MemRead, pc_plus_imm, pc_plus_4, branch_target, pc_sel, ALU_result, rs1, rs2, rd, tempA, tempB, funct3, funct7


MEMORY = data_mem()


def MEM(MEMORY, MemtoReg, RegWrite, MemWrite, RWSel, MemRead, pc_plus_imm,
        pc_plus_4, ALU_result, rs1, rs2, rd, tempA, tempB, funct3, funct7):
    data_out = MEMORY(MemWrite, MemRead, ALU_result, tempB, funct3)
    return MemtoReg, RegWrite, RWSel, pc_plus_imm, pc_plus_4, data_out, tempB, rd


def WB(MemtoReg, RegWrite, RWSel, pc_plus_imm, pc_plus_4, data_out, tempB,
       ex_mem_rd, mem_wb_rd, ex_mem_rs1, ex_mem_rs2, ex_mem_regwrite,
       mem_wb_regwrite):
    mem_out = mux(data_out, tempB, MemtoReg)
    WB_data = mux4(mem_out, pc_plus_4, pc_plus_imm, 0, RWSel)
    forward_a, forward_b = fowardingunit(ex_mem_rs1, ex_mem_rs2, ex_mem_rd,
                                         mem_wb_rd, ex_mem_regwrite,
                                         mem_wb_regwrite)
    return WB_data, RegWrite, forward_a, forward_b