import argparse
import csv

def generate_ratios(max_layers, max_width):
    ratios = []
    for i in range(1, 9):  # 1/8 steps from 1:1 to 1:3 and vice versa
        ratio = i / 8
        if ratio <= 6:  # Ensure the ratio is not greater than 3
            ratios.append(ratio)
    return ratios

def generate_experiments(max_layers, max_width):
    experiments = []
    ratios = generate_ratios(max_layers, max_width)

    for ratio in ratios:
        for layers in range(1, max_layers + 1):
            width = int(max_width * ratio)
            experiments.append((layers, width))

    return experiments

def estimate_vram(layer_count, width, input_size, output_size):
    # Calculate the number of parameters
    # Each layer has (input_size * width) + width parameters (weights + biases)
    # The last layer has (width * output_size) + output_size parameters
    param_count = 0
    for i in range(layer_count):
        if i == 0:
            param_count += (input_size * width) + width
        else:
            param_count += (width * width) + width
    param_count += (width * output_size) + output_size

    # Estimate the VRAM usage
    # Parameters: 4 bytes per parameter (FP32)
    # Activations: Assume the size of the activations is the same as the input size
    vram_usage = param_count * 4 + input_size * 4

    return vram_usage

def write_csv(experiments, filename):
    with open(filename, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['layer_count', 'width', 'vram_usage'])
        for experiment in experiments:
            writer.writerow(experiment)

def main():
    parser = argparse.ArgumentParser(description='Generate a CSV file with a variety of layer counts and widths.')
    parser.add_argument('--max_layers', type=int, default=72, help='Maximum number of layers (default: 4)')
    parser.add_argument('--max_width', type=int, default=(4096), help='Maximum width (default: 2048)')
    parser.add_argument('--output_file', type=str, default='experiments.csv', help='Output CSV file (default: experiments.csv)')
    parser.add_argument('--input_size', type=int, default=64*64*3, help='Input size (default: 224*224*3)')
    parser.add_argument('--output_size', type=int, default=10, help='Output size (default: 10)')
    args = parser.parse_args()

    experiments = generate_experiments(args.max_layers, args.max_width)
    experiments_with_vram = []

    for experiment in experiments:
        layer_count, width = experiment
        vram_usage = estimate_vram(layer_count, width, args.input_size, args.output_size)
        experiments_with_vram.append((layer_count, width, vram_usage))
        print(f'Layer Count: {layer_count}, Width: {width}, Estimated VRAM Usage: {vram_usage} bytes')

    write_csv(experiments_with_vram, args.output_file)
    print(f'Generated {len(experiments_with_vram)} experiments and saved to {args.output_file}')

if __name__ == '__main__':
    main()