import cv2
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import os

class DayLightConverter:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def convert_to_daylight(self, image):
        """夜間画像を昼間のような画像に変換
        Args:
            image: numpy.ndarray (OpenCV形式のBGR画像データ)
        Returns:
            numpy.ndarray: 変換後の画像データ
        """
        if image is None or not isinstance(image, np.ndarray):
            raise ValueError("有効な画像データではありません")

        # 1. 輝度とコントラストの初期調整
        image_lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(image_lab)

        # 輝度チャンネルのヒストグラム解析
        hist = cv2.calcHist([l], [0], None, [256], [0, 256])
        hist_normalized = hist.ravel() / hist.sum()

        # 暗部の割合を計算
        dark_threshold = 100  # 暗部の閾値
        dark_ratio = hist_normalized[:dark_threshold].sum()

        # 輝度補正の強度を自動調整
        gamma = 0.8 if dark_ratio > 0.5 else 1.2
        l_enhanced = np.power(l / 255.0, gamma) * 255.0

        # 2. 適応的なコントラスト強調
        clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(8,8))
        l_enhanced = clahe.apply(l_enhanced.astype(np.uint8))

        # 3. カラーバランスの補正
        a_balanced = cv2.normalize(a, None, 0, 255, cv2.NORM_MINMAX)
        b_balanced = cv2.normalize(b, None, 0, 255, cv2.NORM_MINMAX)

        # 色相の自然な調整（昼光効果）
        a_balanced = cv2.addWeighted(a_balanced, 0.9, np.ones_like(a_balanced) * 128, 0.1, 0)
        b_balanced = cv2.addWeighted(b_balanced, 0.9, np.ones_like(b_balanced) * 128, 0.1, 0)

        # 4. 画像の再構成
        enhanced_lab = cv2.merge([l_enhanced, a_balanced, b_balanced])
        enhanced_bgr = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2BGR)

        # 5. 最終的な色調補正
        # 明るさとコントラストの微調整
        alpha = 1.2  # コントラスト
        beta = 15    # 明るさ
        enhanced_bgr = cv2.convertScaleAbs(enhanced_bgr, alpha=alpha, beta=beta)

        # 6. 自然な見た目のための後処理
        # ホワイトバランスの自動調整
        enhanced_bgr = self._auto_white_balance(enhanced_bgr)

        # エッジ保持ノイズ除去
        enhanced_bgr = cv2.bilateralFilter(enhanced_bgr, d=5, sigmaColor=50, sigmaSpace=50)

        return enhanced_bgr

    def _auto_white_balance(self, image):
        """自動ホワイトバランス調整"""
        result = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        avg_a = np.average(result[:, :, 1])
        avg_b = np.average(result[:, :, 2])

        result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1)
        result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1)

        return cv2.cvtColor(result, cv2.COLOR_LAB2BGR)

def main():
    try:
        converter = DayLightConverter()
        image_path = '1.png'  # 実際の画像パスに変更してください

        # 画像の読み込み
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"画像ファイルの読み込みに失敗しました: {image_path}")

        enhanced_image = converter.convert_to_daylight(image)

        # 出力ディレクトリの確認と作成
        output_dir = 'daylight_converted'
        os.makedirs(output_dir, exist_ok=True)

        output_path = os.path.join(output_dir, 'daylight_' + os.path.basename(image_path))
        cv2.imwrite(output_path, enhanced_image)
        print(f"変換が完了しました。出力ファイル: {output_path}")

    except Exception as e:
        print(f"エラーが発生しました: {e}")

if __name__ == "__main__":
    main()