#!/usr/bin/env python3
# Copyright © Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""
Convert MIOpen Driver Commands to CSV Test Cases

Parses MIOpen driver commands from log files and converts them to CSV format
for CK convolution testing.

Usage:
    python3 miopen_to_csv.py --input miopen_commands.txt --output conv_cases.csv
    python3 miopen_to_csv.py --input miopen_log.txt --output-2d conv_2d.csv --output-3d conv_3d.csv
"""

import argparse
import csv
import re
import os


def parse_miopen_command(command_line):
    """
    Parse MIOpen driver command line into parameter dictionary

    Example input:
    ./bin/MIOpenDriver conv -n 4 -c 3 -H 224 -W 224 -k 64 -y 3 -x 3 -p 1 -q 1 -u 1 -v 1 -l 1 -j 1 -m conv -g 1 -F 1 -t 1

    Returns dict with parsed parameters or None if parsing fails
    """
    if not command_line.strip().startswith("./bin/MIOpenDriver conv"):
        return None

    # Extract parameters using regex
    params = {}

    # Parameter mapping: flag -> description
    # Support both short (-D) and long (--in_d) parameter formats
    param_patterns = {
        "n": r"-n\s+(\d+)",  # batch size
        "c": r"-c\s+(\d+)",  # input channels
        "k": r"-k\s+(\d+)",  # output channels
        "H": r"-H\s+(\d+)",  # input height
        "W": r"-W\s+(\d+)",  # input width
        "D": r"(?:-D|--in_d)\s+(\d+)",  # input depth (3D only) - supports both -D and --in_d
        "y": r"-y\s+(\d+)",  # kernel height
        "x": r"-x\s+(\d+)",  # kernel width
        "z": r"(?:-z|--fil_d)\s+(\d+)",  # kernel depth (3D only) - supports both -z and --fil_d
        "u": r"-u\s+(\d+)",  # stride height
        "v": r"-v\s+(\d+)",  # stride width
        "w": r"(?:-w|--conv_stride_d)\s+(\d+)",  # stride depth (3D only) - supports both -w and --conv_stride_d
        "p": r"-p\s+(\d+)",  # pad height
        "q": r"-q\s+(\d+)",  # pad width
        "s": r"(?:-s|--pad_d)\s+(\d+)",  # pad depth (3D only) - supports both -s and --pad_d
        "l": r"-l\s+(\d+)",  # dilation height
        "j": r"-j\s+(\d+)",  # dilation width
        "r": r"(?:-r|--dilation_d)\s+(\d+)",  # dilation depth (3D only) - supports both -r and --dilation_d
        "g": r"-g\s+(\d+)",  # groups
        "F": r"-F\s+(\d+)",  # direction (1=fwd, 2=bwd_weight, 4=bwd_data)
    }

    for param, pattern in param_patterns.items():
        match = re.search(pattern, command_line)
        if match:
            params[param] = int(match.group(1))

    return params if params else None


def miopen_to_conv_param(miopen_params):
    """
    Convert MIOpen parameters to CK ConvParam format

    Returns dictionary in CSV format or None if conversion fails
    """
    if not miopen_params:
        return None

    # Determine if 2D or 3D convolution
    is_3d = (
        "D" in miopen_params
        or "z" in miopen_params
        or "w" in miopen_params
        or "r" in miopen_params
        or "s" in miopen_params
    )

    # Extract basic parameters with defaults
    ndim = 3 if is_3d else 2
    groups = miopen_params.get("g", 1)
    batch_size = miopen_params.get("n", 1)
    # MIOpen uses total channels (C*G), CK uses channels per group
    out_channels_total = miopen_params.get("k", 64)
    in_channels_total = miopen_params.get("c", 3)
    out_channels = out_channels_total // groups  # CK format: channels per group
    in_channels = in_channels_total // groups  # CK format: channels per group

    if is_3d:
        # 3D convolution
        kernel_d = miopen_params.get("z", 3)
        kernel_h = miopen_params.get("y", 3)
        kernel_w = miopen_params.get("x", 3)

        input_d = miopen_params.get("D", 16)
        input_h = miopen_params.get("H", 32)
        input_w = miopen_params.get("W", 32)

        stride_d = miopen_params.get("w", 1)
        stride_h = miopen_params.get("u", 1)
        stride_w = miopen_params.get("v", 1)

        dilation_d = miopen_params.get("r", 1)
        dilation_h = miopen_params.get("l", 1)
        dilation_w = miopen_params.get("j", 1)

        pad_d = miopen_params.get("s", 0)
        pad_h = miopen_params.get("p", 0)
        pad_w = miopen_params.get("q", 0)

        # Calculate output dimensions
        output_d = (
            input_d + 2 * pad_d - dilation_d * (kernel_d - 1) - 1
        ) // stride_d + 1
        output_h = (
            input_h + 2 * pad_h - dilation_h * (kernel_h - 1) - 1
        ) // stride_h + 1
        output_w = (
            input_w + 2 * pad_w - dilation_w * (kernel_w - 1) - 1
        ) // stride_w + 1

        # Skip invalid configurations
        if output_d <= 0 or output_h <= 0 or output_w <= 0:
            return None

        direction = miopen_params.get("F", 1)  # 1=fwd, 2=bwd_weight, 4=bwd_data
        direction_name = {1: "fwd", 2: "bwd_weight", 4: "bwd_data"}.get(
            direction, "fwd"
        )

        return {
            "NDim": ndim,
            "Groups": groups,
            "BatchSize": batch_size,
            "OutChannels": out_channels,
            "InChannels": in_channels,
            "KernelD": kernel_d,
            "KernelH": kernel_h,
            "KernelW": kernel_w,
            "InputD": input_d,
            "InputH": input_h,
            "InputW": input_w,
            "OutputD": output_d,
            "OutputH": output_h,
            "OutputW": output_w,
            "StrideD": stride_d,
            "StrideH": stride_h,
            "StrideW": stride_w,
            "DilationD": dilation_d,
            "DilationH": dilation_h,
            "DilationW": dilation_w,
            "LeftPadD": pad_d,
            "LeftPadH": pad_h,
            "LeftPadW": pad_w,
            "RightPadD": pad_d,
            "RightPadH": pad_h,
            "RightPadW": pad_w,
            "TestName": f"MIOpen_3D_{direction_name}",
        }

    else:
        # 2D convolution
        kernel_h = miopen_params.get("y", 3)
        kernel_w = miopen_params.get("x", 3)

        input_h = miopen_params.get("H", 32)
        input_w = miopen_params.get("W", 32)

        stride_h = miopen_params.get("u", 1)
        stride_w = miopen_params.get("v", 1)

        dilation_h = miopen_params.get("l", 1)
        dilation_w = miopen_params.get("j", 1)

        pad_h = miopen_params.get("p", 0)
        pad_w = miopen_params.get("q", 0)

        # Calculate output dimensions
        output_h = (
            input_h + 2 * pad_h - dilation_h * (kernel_h - 1) - 1
        ) // stride_h + 1
        output_w = (
            input_w + 2 * pad_w - dilation_w * (kernel_w - 1) - 1
        ) // stride_w + 1

        # Skip invalid configurations
        if output_h <= 0 or output_w <= 0:
            return None

        direction = miopen_params.get("F", 1)
        direction_name = {1: "fwd", 2: "bwd_weight", 4: "bwd_data"}.get(
            direction, "fwd"
        )

        return {
            "NDim": ndim,
            "Groups": groups,
            "BatchSize": batch_size,
            "OutChannels": out_channels,
            "InChannels": in_channels,
            "KernelH": kernel_h,
            "KernelW": kernel_w,
            "InputH": input_h,
            "InputW": input_w,
            "OutputH": output_h,
            "OutputW": output_w,
            "StrideH": stride_h,
            "StrideW": stride_w,
            "DilationH": dilation_h,
            "DilationW": dilation_w,
            "LeftPadH": pad_h,
            "LeftPadW": pad_w,
            "RightPadH": pad_h,
            "RightPadW": pad_w,
            "TestName": f"MIOpen_2D_{direction_name}",
        }


def write_csv_cases(test_cases, output_file, ndim):
    """Write test cases to CSV file"""
    if not test_cases:
        print(f"No {ndim}D test cases to write")
        return

    print(f"Writing {len(test_cases)} {ndim}D test cases to {output_file}")

    # Define CSV headers based on dimension
    if ndim == 2:
        headers = [
            "NDim",
            "Groups",
            "BatchSize",
            "OutChannels",
            "InChannels",
            "KernelH",
            "KernelW",
            "InputH",
            "InputW",
            "OutputH",
            "OutputW",
            "StrideH",
            "StrideW",
            "DilationH",
            "DilationW",
            "LeftPadH",
            "LeftPadW",
            "RightPadH",
            "RightPadW",
            "TestName",
        ]
    else:  # 3D
        headers = [
            "NDim",
            "Groups",
            "BatchSize",
            "OutChannels",
            "InChannels",
            "KernelD",
            "KernelH",
            "KernelW",
            "InputD",
            "InputH",
            "InputW",
            "OutputD",
            "OutputH",
            "OutputW",
            "StrideD",
            "StrideH",
            "StrideW",
            "DilationD",
            "DilationH",
            "DilationW",
            "LeftPadD",
            "LeftPadH",
            "LeftPadW",
            "RightPadD",
            "RightPadH",
            "RightPadW",
            "TestName",
        ]

    with open(output_file, "w", newline="") as csvfile:
        # Write header comment
        csvfile.write(f"# {ndim}D Convolution Test Cases from MIOpen Commands\n")
        csvfile.write(f"# Generated {len(test_cases)} test cases\n")

        writer = csv.DictWriter(csvfile, fieldnames=headers)
        writer.writeheader()

        for test_case in test_cases:
            # Only write fields that exist in headers
            filtered_case = {k: v for k, v in test_case.items() if k in headers}
            writer.writerow(filtered_case)


def main():
    parser = argparse.ArgumentParser(
        description="Convert MIOpen commands to CSV test cases"
    )

    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="Input file with MIOpen driver commands",
    )
    parser.add_argument(
        "--output", type=str, help="Output CSV file (for mixed 2D/3D cases)"
    )
    parser.add_argument(
        "--output-2d",
        type=str,
        default="miopen_conv_2d.csv",
        help="Output CSV file for 2D cases",
    )
    parser.add_argument(
        "--output-3d",
        type=str,
        default="miopen_conv_3d.csv",
        help="Output CSV file for 3D cases",
    )
    parser.add_argument(
        "--filter-duplicates", action="store_true", help="Remove duplicate test cases"
    )
    parser.add_argument(
        "--model-name",
        type=str,
        default="MIOpen",
        help="Model name to use in test case names (default: MIOpen)",
    )

    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"ERROR: Input file not found: {args.input}")
        return 1

    print(f"Parsing MIOpen commands from {args.input}...")

    test_cases_2d = []
    test_cases_3d = []
    total_lines = 0
    parsed_lines = 0

    with open(args.input, "r") as f:
        for line_num, line in enumerate(f, 1):
            total_lines += 1
            line = line.strip()

            # Skip empty lines and non-MIOpen commands
            # Handle both direct commands and logged commands with MIOpen prefix
            if not line:
                continue

            # Extract the actual MIOpenDriver command from logged format
            if "MIOpenDriver conv" in line:
                # Extract command after finding MIOpenDriver
                command_start = line.find("./bin/MIOpenDriver conv")
                if command_start != -1:
                    line = line[command_start:]
                else:
                    # Handle cases where path might be different - create standard format
                    driver_start = line.find("MIOpenDriver conv")
                    if driver_start != -1:
                        line = "./bin/" + line[driver_start:]
                    else:
                        continue
            elif not line.startswith("./bin/MIOpenDriver conv"):
                continue

            try:
                # Parse MIOpen command
                miopen_params = parse_miopen_command(line)
                if not miopen_params:
                    continue

                # Convert to ConvParam format
                conv_param = miopen_to_conv_param(miopen_params)
                if not conv_param:
                    continue

                # Add model name to test name
                conv_param["TestName"] = f"{args.model_name}_{conv_param['NDim']}D_fwd"

                # Separate 2D and 3D cases
                if conv_param["NDim"] == 2:
                    test_cases_2d.append(conv_param)
                else:
                    test_cases_3d.append(conv_param)

                parsed_lines += 1

            except Exception as e:
                print(f"WARNING: Failed to parse line {line_num}: {e}")
                continue

    print(f"Processed {total_lines} lines, parsed {parsed_lines} commands")
    print(f"Found {len(test_cases_2d)} 2D cases, {len(test_cases_3d)} 3D cases")

    # Remove duplicates if requested
    if args.filter_duplicates:
        # Simple duplicate removal based on key parameters
        def make_key(case):
            if case["NDim"] == 2:
                return (
                    case["Groups"],
                    case["BatchSize"],
                    case["OutChannels"],
                    case["InChannels"],
                    case["KernelH"],
                    case["KernelW"],
                    case["InputH"],
                    case["InputW"],
                    case["StrideH"],
                    case["StrideW"],
                )
            else:
                return (
                    case["Groups"],
                    case["BatchSize"],
                    case["OutChannels"],
                    case["InChannels"],
                    case["KernelD"],
                    case["KernelH"],
                    case["KernelW"],
                    case["InputD"],
                    case["InputH"],
                    case["InputW"],
                    case["StrideD"],
                    case["StrideH"],
                    case["StrideW"],
                )

        seen_2d = set()
        unique_2d = []
        for case in test_cases_2d:
            key = make_key(case)
            if key not in seen_2d:
                seen_2d.add(key)
                unique_2d.append(case)

        seen_3d = set()
        unique_3d = []
        for case in test_cases_3d:
            key = make_key(case)
            if key not in seen_3d:
                seen_3d.add(key)
                unique_3d.append(case)

        print(
            f"After deduplication: {len(unique_2d)} 2D cases, {len(unique_3d)} 3D cases"
        )
        test_cases_2d = unique_2d
        test_cases_3d = unique_3d

    # Write output files
    if args.output:
        # Write mixed cases to single file
        all_cases = test_cases_2d + test_cases_3d
        if all_cases:
            print(f"Writing {len(all_cases)} total cases to {args.output}")
            # Use 2D headers for mixed file, extend as needed
            mixed_headers = [
                "NDim",
                "Groups",
                "BatchSize",
                "OutChannels",
                "InChannels",
                "KernelH",
                "KernelW",
                "InputH",
                "InputW",
                "OutputH",
                "OutputW",
                "StrideH",
                "StrideW",
                "DilationH",
                "DilationW",
                "LeftPadH",
                "LeftPadW",
                "RightPadH",
                "RightPadW",
                "TestName",
            ]

            with open(args.output, "w", newline="") as csvfile:
                csvfile.write(
                    "# Mixed 2D/3D Convolution Test Cases from MIOpen Commands\n"
                )
                writer = csv.DictWriter(
                    csvfile, fieldnames=mixed_headers, extrasaction="ignore"
                )
                writer.writeheader()
                for case in all_cases:
                    writer.writerow(case)
    else:
        # Write separate files for 2D and 3D
        if test_cases_2d:
            write_csv_cases(test_cases_2d, args.output_2d, 2)

        if test_cases_3d:
            write_csv_cases(test_cases_3d, args.output_3d, 3)

    print("Conversion completed!")
    return 0


if __name__ == "__main__":
    exit(main())
