import os
import cv2
import pandas as pd
import re

class ImageProcessor:
    def __init__(self):
        self.output_dir = "images-fix"
        os.makedirs(self.output_dir, exist_ok=True)
        # pos番号の変換マップを定義
        self.pos_map = {
            '3': '4',
            '4': '3',
            '7': '8',
            '8': '7'
        }
        
    def get_new_filename(self, original_filename):
        """ファイル名のpos番号を必要に応じて変更する"""
        pattern = r'(.*pos)(\d+)(.*)'
        match = re.match(pattern, original_filename)
        if match:
            pos_num = match.group(2)
            if pos_num in self.pos_map:
                return re.sub(pattern, f'\\g<1>{self.pos_map[pos_num]}\\g<3>', original_filename)
        return original_filename
    
    def check_duplicate_filename(self, filename):
        """出力ディレクトリで重複をチェック"""
        return os.path.exists(os.path.join(self.output_dir, filename))
    
    def process_folder(self, folder_path):
        results_csv = os.path.join(folder_path, "results.csv")
        images_folder = os.path.join(folder_path, "images")
        
        if not os.path.exists(results_csv):
            print(f"結果ファイルが見つかりません: {results_csv}")
            return
        
        df = pd.read_csv(results_csv)
        
        for _, row in df.iterrows():
            image_file_name = row["image_file_name"]
            image_path = os.path.join(images_folder, image_file_name)
            
            if not os.path.exists(image_path):
                print(f"画像ファイルが見つかりません: {image_path}")
                continue
                
            image = cv2.imread(image_path)
            if image is None:
                print(f"画像の読み込みに失敗しました: {image_path}")
                continue
            
            # 新しいファイル名を取得
            new_filename = self.get_new_filename(image_file_name)
            output_path = os.path.join(self.output_dir, new_filename)
            
            # ファイルを保存
            cv2.imwrite(output_path, image)
            print(f"保存完了: {output_path}")

def get_numeric_folders(base_dir="data"):
    """
    dataディレクトリ���の数字名フォルダを取得する
    """
    folders = []
    if os.path.exists(base_dir):
        for item in os.listdir(base_dir):
            if os.path.isdir(os.path.join(base_dir, item)) and item.isdigit():
                folders.append(item)
    return sorted(folders, key=lambda x: int(x))

def main():
    processor = ImageProcessor()
    
    # dataディレクトリ内の数字フォルダを取得
    numeric_folders = get_numeric_folders()
    if not numeric_folders:
        print("処理対象のフォルダが見つかりません")
        return
    
    print("処理するフォルダを選択してください:")
    for i, folder in enumerate(numeric_folders, 1):
        print(f"{i}: {folder}")
    print("0: 全てのフォルダを処理")
    
    while True:
        try:
            choice = input("番号を入力してください: ").strip()
            if not choice:
                continue
                
            choice = int(choice)
            if choice == 0:
                # 全てのフォルダを処理
                folders_to_process = numeric_folders
                break
            elif 1 <= choice <= len(numeric_folders):
                # 選択されたフォルダのみ処理
                folders_to_process = [numeric_folders[choice-1]]
                break
            else:
                print(f"1から{len(numeric_folders)}までの数字、または0を入力してください")
        except ValueError:
            print("有効な数字を入力してください")
    
    # 選択されたフォルダを処理
    for folder in folders_to_process:
        folder_path = os.path.join("data", folder.zfill(2))
        print(f"\n=== フォルダ {folder} の処理を開始 ===")
        processor.process_folder(folder_path)
    
    print("\n処理が完了しました")

if __name__ == "__main__":
    main()