# Copyright 2022-2026 XMOS LIMITED.
# This Software is subject to the terms of the XMOS Public Licence: Version 1.
import argparse
import numpy as np
from pathlib import Path

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", nargs='?', default='2 2 2 10 5', help="Build configuration in '<threads> <num_y_channels> <num_x_channels> <num_main_filter_phases> <num_shadow_filter_phases>' format. Default '2 2 2 10 5'")
    parser.add_argument("--out-dir", nargs='?', default='.', help="output directory to generate files in. Default .")
    args = parser.parse_args()
    return args

#Schedule num_tasks tasks, each having num_channels channels over num_threads threads
def gen_task_distribution(num_threads, num_tasks, num_channels):
    #Scheduling logic - T0C0, T1C0, T2C0, T0C1, T1C1, T2C2 and so on
    if num_channels:
        remaining_jobs = num_tasks * num_channels
    else:
        remaining_jobs = num_tasks;
    pas = 0
    tdist = np.zeros((num_threads,1), dtype=np.int32)
    thread = 0

    task_index = 0
    ch_index = 0
    while(remaining_jobs):
        #bit 16:9 task_index, bit8:1 ch_index, bit0, is_active
        tdist[thread, pas] = (1 | (((ch_index & 0xff) << 1) | ((task_index & 0xff) << 9)))
        task_index = task_index + 1
        if(task_index == num_tasks):
            task_index = 0
            if num_channels:
                ch_index = (ch_index + 1) % num_channels

        thread += 1
        remaining_jobs -= 1
        if(thread == num_threads and remaining_jobs): #add another pass
            thread = 0
            pas += 1
            append_sch = np.zeros((num_threads, 1), dtype=np.int32)
            tdist = np.hstack((tdist, append_sch))
    pas += 1
    return tdist

def print_task_distribution(tdist, num_tasks, num_channels, fp_c):
    num_threads = tdist.shape[0]
    num_passes = tdist.shape[-1]
    if num_channels:
        passes_c_str = f"passes_for_{num_tasks}_tasks_and_channels"
        par_str = f"par_{num_tasks}_tasks_and_channels"
    else:
        passes_c_str = f"passes_for_{num_tasks}_tasks"
        par_str = f"par_{num_tasks}_tasks"

    fp_c.write(f"\t.{passes_c_str} = {num_passes},\n")
    fp_c.write((f"\t.{par_str} = " "{\n"))

    for t in range(num_threads):
        thread_str = "\t\t{"
        for p in range(num_passes):
            task = (tdist[t][p] >> 9);
            ch = (tdist[t][p] >> 1) & 0xff;
            active = tdist[t][p] & 1;
            if num_channels:
                ph_str = "{" + f"{task}, {ch}, {active}" "}"
            else:
                ph_str = "{" + f"{task}, {active}" "}"
            thread_str = thread_str + ph_str + ','
        thread_str += "},"
        fp_c.write(thread_str+'\n')
    fp_c.write("\t},\n")


def create_task_distribution():
    args = parse_arguments()
    print('config = ',args.config)
    print('out-dir = ',args.out_dir)
    conf = args.config.split(' ')
    threads = conf[0]
    max_y_channels = conf[1]
    max_x_channels = conf[2]
    main_filter_phases = conf[3]
    shadow_filter_phases = conf[4]
    print(f"AEC task distribution generation configured for {threads} threads, max {max_y_channels} y channels, max {max_x_channels} x channels, {main_filter_phases} main_filter_phases, {shadow_filter_phases} shadow_filter_phases")
    autogen_message = '/* Do not edit, autogenerated */ '

    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
    cfg_h_file_name = Path(args.out_dir) / "aec_conf.h"
    c_file_name = Path(args.out_dir) / "aec_task_distribution.c"

    fp_cfg_h = open(cfg_h_file_name, 'w')
    fp_c = open(c_file_name, 'w')

    with open(cfg_h_file_name, 'w') as fp_cfg_h, open(c_file_name, 'w') as fp_c:
        fp_cfg_h.write( '#ifndef aec_config_h_\n')
        fp_cfg_h.write( '#define aec_config_h_\n')
        # Distribute multiple tasks across multiple channels on different cores
        num_channels = max(int(max_y_channels), int(max_x_channels))
        num_threads = int(threads)
        fp_cfg_h.write(autogen_message + '\n')
        fp_cfg_h.write(f"#define AEC_MAX_Y_CHANNELS   ({max_y_channels})\n")
        fp_cfg_h.write(f"#define AEC_MAX_X_CHANNELS   ({max_x_channels})\n")
        fp_cfg_h.write(f"#define AEC_MAIN_FILTER_PHASES    ({main_filter_phases})\n")
        fp_cfg_h.write(f"#define AEC_SHADOW_FILTER_PHASES    ({shadow_filter_phases})\n")

        fp_c.write( autogen_message + '\n')
        fp_c.write( '#include "aec_schedule.h"\n\n')
        fp_c.write( 'aec_task_distribution_t tdist = {\n')
        fp_c.write( f'\t.thread_count = {num_threads},\n')
        #Schedule 3 tasks, num_channels channels over num_threads threads
        tdist = gen_task_distribution(num_threads, 3, num_channels)
        print_task_distribution(tdist, 3, num_channels, fp_c);

        #Schedule 2 tasks, num_channels channels over num_threads threads
        tdist = gen_task_distribution(num_threads, 2, num_channels)
        print_task_distribution(tdist, 2, num_channels, fp_c)

        #Schedule 1 tasks, num_channels channels over num_threads threads
        tdist = gen_task_distribution(num_threads, 1, num_channels)
        print_task_distribution(tdist, 1, num_channels, fp_c)

        #Schedule multiple tasks on different cores
        #Schedule 3 tasks
        tdist = gen_task_distribution(num_threads, 3, 0)
        print_task_distribution(tdist, 3, 0, fp_c)

        #Schedule 2 tasks
        tdist = gen_task_distribution(num_threads, 2, 0)
        print_task_distribution(tdist, 2, 0, fp_c)

        fp_cfg_h.write( '#endif /* aec_config_h_ */')
        fp_c.write( '};\n\n')

if __name__ == "__main__":
    create_task_distribution()
