from bs4 import BeautifulSoup
from templates import *

class taskdata:
    def __init__(self, delay, count, interval, type, task_id, dest_id, req_id):
        self.delay = delay
        self.count = count
        self.interval = interval
        self.type = type
        self.task_id = task_id
        self.dest_id = dest_id
        self.req_id = req_id

class xmlformater:
    def __init__(self):
        self.tasks = DATA_XML_START
        self.xml_maps = ""
        self.curr_dest_task = 1
        self.count_wk = {}
        self.count_if = {}
        self.repetitions_wk = {}
        self.repetitions_if = {}
        self.maps = {0:PE_GB}

    def read_tasks(self, data_file):
        with open(data_file, 'r') as f:
            data = f.read()
        Bs_data = BeautifulSoup(data, "xml")
        tasks = Bs_data.find_all('task')
        for task in tasks:
            id = task.get('id')
            if id is not None:
                dests = task.find_all('destination')
                if dests:
                    for dest in dests:
                        count = dest.get('count')
                        type = dest.get('type')
                        aa = dest.get('type')
                        dest_task  = dests[0].find('task').get('value')


    def generate_task(self, num_weights, size_weights, size_input_features, size_output):
        dests = ""
        dest_task_start = 1
        task_id = 0
        dest_id = 0
        delay = DELAY
        interval = INTERVAL

        total_wk_values = num_weights * size_weights[0] * size_weights[1] * size_weights[2]
        wk_per_chiplet = total_wk_values // NUM_CHIPLETS
        remainder_wk = total_wk_values % NUM_CHIPLETS  # Extra values to distribute

        # Compute IF Distribution (Input Feature Maps)
        total_if_chunks = (size_output[0] * size_output[1])  # 112 * 112 = 12544
        if_per_chiplet = total_if_chunks // NUM_CHIPLETS
        remainder_if = total_if_chunks % NUM_CHIPLETS  # Extra values to distribute


        for i in range(0,NUM_PE): # WK
            count = wk_per_chiplet + (1 if i < remainder_wk else 0)  # Distribute remainder across chiplets
            # repetitions = num_weights / NUM_PE -1
            # self.repetitions_wk[i] = 0
            # if i < repetitions:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2]*(repetitions+1))
            #     self.repetitions_wk[i] = repetitions
            # elif repetitions > 0:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2]*repetitions)
            # else:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2])
            dest_task = dest_task_start + i
            dests += DEST.format(dest_id,delay,delay,interval,interval,count,count,TYPE_WK,dest_task) #!
            self.count_wk[i] = count
            self.maps[self.curr_dest_task] = START_PE + i
            dest_id += 1
            self.curr_dest_task += 1
            

        for i in range(0,NUM_CHIPLETS): # IF
            count = if_per_chiplet + (1 if i < remainder_if else 0)
            count *= size_weights[0] * size_weights[1] * size_weights[2]  # Multiply by 27 (3x3x3 input chunk size)
            # repetitions = (size_output[0]*size_output[1]*size_output[2] - NUM_CHIPLETS*NUM_PE) / NUM_CHIPLETS
            # self.repetitions_if[i] = 0
            # if i < repetitions:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2]*(repetitions+1))
            #     self.repetitions_if[i] = repetitions
            # elif repetitions > 0:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2]*repetitions)
            # else:
            #     count = int(size_weights[0]*size_weights[1]*size_weights[2])
            dest_task = dest_task_start + NUM_PE * i
            dests += DEST.format(dest_id,delay,delay,interval,interval,count,count,TYPE_IF,dest_task)
            self.count_if[i] = count
            self.maps[self.curr_dest_task] = START_PE + i*NUM_PE
            dest_id += 1
            self.curr_dest_task += 1
    
        self.tasks += TASK.format(task_id, dests, "")
        self.curr_dest_task += 1


    def generate_receiving_tasks(self,size_output):
        total_outputs = size_output[0] * size_output[1] * size_output[2]  # 112 * 112 * 32 = 401408
        outputs_per_pe = total_outputs // (NUM_CHIPLETS * NUM_PE)  # 401408 / 64 = 6272
        remainder_outputs = total_outputs % (NUM_CHIPLETS * NUM_PE)  # Extra outputs to distribute
        sum_count = 0
        last_task_id = 1 + NUM_PE * NUM_CHIPLETS
        for i in range(0, NUM_CHIPLETS):
            for j in range(0, NUM_PE):
                reqs = ""
                req_id = 0
                reqs += REQ.format(req_id, TYPE_IF, 0, self.count_if[i], self.count_if[i])
                req_id += 1
                reqs += REQ.format(req_id, TYPE_WK, 0, self.count_wk[j], self.count_wk[j])

                # count = int(self.repetitions_if[i]+self.repetitions_wk[j]+1)
                count = outputs_per_pe + (1 if (i * NUM_PE + j) < remainder_outputs else 0)  # Each PE generates 6272 outputs
                sum_count += count
                dest = DEST.format(0,DELAY,DELAY,INTERVAL,INTERVAL,count,count,TYPE_OF,last_task_id)

                task_id = 1+i*NUM_PE+j
                self.tasks += TASK.format(task_id,dest,reqs)
    
                self.maps[task_id] = START_PE +i*NUM_PE+j

        req = REQ.format(0, TYPE_OF, 0, sum_count, sum_count)
        self.tasks += TASK.format(last_task_id, "", req)
        self.maps[last_task_id] = PE_GB


    def generate_maps(self):
        binds = ""
        for i in range(0, len(self.maps)):
            binds += BIND.format(i,self.maps[i])
        self.xml_maps = MAP.format(binds)

    def write_xml_file(self):
        f = open("./my_data.xml", "w")
        f.write("")
        f.close()

        self.tasks += DATA_XML_END

        f = open("./my_data.xml", "a")
        f.write(self.tasks)
        f.close()


        f = open("./my_map.xml", "w")
        f.write("")
        f.close()

        f = open("./my_map.xml", "a")
        f.write(self.xml_maps)
        f.close()