import pandas as pd
import os
import sys
from colorama import Fore, Style, init

def create_damage_files(base_path, output_base, prefix):
    # Initialize colorama
    init(autoreset=True)
    
    # Generate column labels based on expected duplication in input files
    columns = ['Real'] + [f'Real.{i}' for i in range(1, 30)]  # Explicitly setting column names

    sensor_end_map = {1: 'Real.25', 2: 'Real.26', 3: 'Real.27', 4: 'Real.28', 5: 'Real.29'}

    # Define the damage scenarios and the corresponding original file indices
    damage_scenarios = {
        1: range(1, 6),  # Damage 1 files from zzzAD1.csv to zzzAD5.csv
        2: range(6, 11),  # Damage 2 files from zzzAD6.csv to zzzAD10.csv
        3: range(11, 16), # Damage 3 files from zzzAD11.csv to zzzAD15.csvs
        4: range(16, 21), # Damage 4 files from zzzAD16.csv to zzzAD20.csv
        5: range(21, 26),  # Damage 5 files from zzzAD21.csv to zzzAD25.csv
        6: range(26, 31)  # Damage 6 files from zzzAD26.csv to zzzAD30.csv
    }
    damage_pad = len(str(len(damage_scenarios)))
    test_pad = len(str(30))

    for damage, files in damage_scenarios.items():
        for i, file_index in enumerate(files, start=1):
            # Load original data file
            file_path = os.path.join(base_path, f'zzz{prefix}D{file_index}.TXT')
            df = pd.read_csv(file_path, sep='\t', skiprows=10)  # Read with explicit column names

            top_sensor = columns[i-1]
            print(top_sensor, type(top_sensor))
            output_file_1 = os.path.join(output_base, f'DAMAGE_{damage}', f'DAMAGE{damage}_TEST{i}_01.csv')
            print(f"Creating {output_file_1} from taking zzz{prefix}D{file_index}.TXT")
            print("Taking datetime column on index 0...")
            print(f"Taking `{top_sensor}`...")
            df[['Time', top_sensor]].to_csv(output_file_1, index=False)
            print(Fore.GREEN + "Done")

            bottom_sensor = sensor_end_map[i]
            output_file_2 = os.path.join(output_base, f'DAMAGE_{damage}', f'DAMAGE{damage}_TEST{i}_02.csv')
            print(f"Creating {output_file_2} from taking zzz{prefix}D{file_index}.TXT")
            print("Taking datetime column on index 0...")
            print(f"Taking `{bottom_sensor}`...")
            df[['Time', bottom_sensor]].to_csv(output_file_2, index=False)
            print(Fore.GREEN + "Done")
            print("---")

def main():
    if len(sys.argv) < 2:
        print("Usage: python convert.py <path_to_csv_files>")
        sys.exit(1)

    base_path = sys.argv[1]
    output_base = sys.argv[2]
    prefix = sys.argv[3]  # Define output directory

    # Create output folders if they don't exist
    for i in range(1, 5):
        os.makedirs(os.path.join(output_base, f'DAMAGE_{i}'), exist_ok=True)

    create_damage_files(base_path, output_base, prefix)
    print(Fore.YELLOW + Style.BRIGHT + "All files have been created successfully.")

if __name__ == "__main__":
    main()