import io
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from flask import Blueprint, jsonify, request

bp = Blueprint('water_body_segmentation', __name__, url_prefix='/api/water_body_segmentation')

# DeepLabV3 모델 (나중에 실제 모델로 교체)
# from app.AI_modules.DeepLabV3 import create_model, load_model
deeplabv3_model = None


def generate_satellite_image(lat, lng):
    """위성 이미지 시뮬레이션 생성 (INPUT)"""
    fig, ax = plt.subplots(figsize=(8, 8))

    np.random.seed(int((lat + lng) * 1000) % 10000)

    # 배경 (녹지/농경지)
    terrain = np.random.rand(100, 100) * 0.3 + 0.3

    # 도로/경계선 추가
    for i in range(5):
        start = np.random.randint(0, 100)
        terrain[start:start+2, :] = 0.9
        terrain[:, start:start+2] = 0.9

    # 물 영역 생성 (불규칙한 형태)
    y, x = np.ogrid[:100, :100]

    # 주요 수역 1 (호수/저수지)
    center1_x, center1_y = 35 + np.random.randint(-5, 5), 50 + np.random.randint(-5, 5)
    dist1 = np.sqrt((x - center1_x)**2 + (y - center1_y)**2)
    water_mask1 = dist1 < (15 + np.random.rand(100, 100) * 5)

    # 주요 수역 2 (강/하천)
    river_path = 60 + np.sin(np.linspace(0, 4*np.pi, 100)) * 10
    for i in range(100):
        width = 3 + int(np.random.rand() * 3)
        river_col = int(river_path[i])
        if 0 <= river_col < 100:
            terrain[i, max(0, river_col-width):min(100, river_col+width)] = 0.2

    terrain[water_mask1] = 0.15 + np.random.rand(np.sum(water_mask1)) * 0.1

    # 녹색 채널 강조 (위성 이미지처럼)
    rgb_image = np.zeros((100, 100, 3))
    rgb_image[:, :, 0] = terrain * 0.4 + 0.1  # R
    rgb_image[:, :, 1] = terrain * 0.8 + 0.1  # G (녹지 강조)
    rgb_image[:, :, 2] = terrain * 0.3 + 0.1  # B

    # 물 영역은 파란색으로
    water_full_mask = (terrain < 0.3)
    rgb_image[water_full_mask, 0] = 0.1
    rgb_image[water_full_mask, 1] = 0.3
    rgb_image[water_full_mask, 2] = 0.5

    ax.imshow(rgb_image, extent=[lng-0.05, lng+0.05, lat-0.05, lat+0.05])
    ax.set_xlabel(f'경도 (Longitude)', fontsize=10)
    ax.set_ylabel(f'위도 (Latitude)', fontsize=10)
    ax.set_title(f'위성 이미지 - ({lat:.4f}, {lng:.4f})', fontsize=12)

    plt.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64, water_full_mask


def generate_segmentation_image(lat, lng, water_mask):
    """Water Body 세그멘테이션 결과 생성 (OUTPUT)"""
    fig, ax = plt.subplots(figsize=(8, 8))

    np.random.seed(int((lat + lng) * 1000) % 10000)

    # 배경 (원본 이미지 흐리게)
    terrain = np.random.rand(100, 100) * 0.3 + 0.3

    # 도로/경계선
    for i in range(5):
        start = np.random.randint(0, 100)
        terrain[start:start+2, :] = 0.9
        terrain[:, start:start+2] = 0.9

    # RGB 이미지 생성 (흐린 배경)
    rgb_image = np.zeros((100, 100, 3))
    rgb_image[:, :, 0] = terrain * 0.5 + 0.3
    rgb_image[:, :, 1] = terrain * 0.6 + 0.3
    rgb_image[:, :, 2] = terrain * 0.5 + 0.3

    # Water Body를 밝은 파란색/보라색으로 하이라이트
    rgb_image[water_mask, 0] = 0.4
    rgb_image[water_mask, 1] = 0.5
    rgb_image[water_mask, 2] = 0.95

    # 경계선 추가
    from scipy import ndimage
    edges = ndimage.binary_dilation(water_mask) & ~water_mask
    rgb_image[edges, 0] = 1.0
    rgb_image[edges, 1] = 1.0
    rgb_image[edges, 2] = 1.0

    ax.imshow(rgb_image, extent=[lng-0.05, lng+0.05, lat-0.05, lat+0.05])
    ax.set_xlabel(f'경도 (Longitude)', fontsize=10)
    ax.set_ylabel(f'위도 (Latitude)', fontsize=10)
    ax.set_title(f'Water Body 추출 결과', fontsize=12)

    # 범례 추가
    from matplotlib.patches import Patch
    legend_elements = [Patch(facecolor='#6680f2', edgecolor='white', label='Water Body')]
    ax.legend(handles=legend_elements, loc='upper right')

    plt.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    # 면적 계산 (시뮬레이션)
    water_pixel_count = np.sum(water_mask)
    total_pixels = water_mask.size
    water_ratio = (water_pixel_count / total_pixels) * 100

    return image_base64, water_pixel_count, water_ratio


@bp.route('/extract', methods=['GET'])
def extract():
    """Water Body 추출 결과 반환"""
    try:
        lat = float(request.args.get('lat', 34.6047))
        lng = float(request.args.get('lng', 127.2855))
        radius = float(request.args.get('radius', 5))

        # 입력 이미지 생성
        input_image, water_mask = generate_satellite_image(lat, lng)

        # 출력 이미지 생성
        output_image, water_pixels, water_ratio = generate_segmentation_image(lat, lng, water_mask)

        # 면적 계산 (시뮬레이션: 1 pixel = 약 100m²)
        area_m2 = water_pixels * 100
        water_body_count = np.random.randint(2, 8)

        return jsonify({
            'status': 'success',
            'message': 'Water Body 추출 완료',
            'input_image': input_image,
            'output_image': output_image,
            'data': {
                'area': f'{area_m2:,} ㎡',
                'ratio': f'{water_ratio:.1f}%',
                'count': water_body_count,
                'lat': lat,
                'lng': lng
            }
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'input_image': None,
            'output_image': None
        })


def simulate_water_segmentation(image_array):
    """
    Water Body 세그멘테이션 시뮬레이션
    (실제 DeepLabV3 모델로 교체 예정)
    """
    height, width = image_array.shape[:2]

    # 파란색 계열 검출 (간단한 색상 기반 검출)
    if len(image_array.shape) == 3:
        # RGB -> HSV 변환 대신 간단한 파란색 검출
        r = image_array[:, :, 0].astype(float)
        g = image_array[:, :, 1].astype(float)
        b = image_array[:, :, 2].astype(float)

        # 파란색/청록색 영역 검출
        blue_ratio = b / (r + g + b + 1e-6)
        green_blue_diff = np.abs(g - b) / 255.0

        # 물로 추정되는 영역 (파란색 비율이 높고, 전체적으로 어두운 영역)
        brightness = (r + g + b) / 3.0
        water_mask = (blue_ratio > 0.35) | ((brightness < 100) & (b > r))

        # 노이즈 제거를 위한 morphological operation 시뮬레이션
        from scipy import ndimage
        water_mask = ndimage.binary_opening(water_mask, iterations=2)
        water_mask = ndimage.binary_closing(water_mask, iterations=2)
    else:
        # 그레이스케일 이미지
        water_mask = image_array < 100

    return water_mask


def generate_prediction_overlay(original_image, water_mask):
    """
    원본 이미지에 Water Body 예측 결과 오버레이
    """
    fig, ax = plt.subplots(figsize=(8, 8))

    # 원본 이미지 표시
    ax.imshow(original_image)

    # Water Body 마스크를 반투명 파란색으로 오버레이
    overlay = np.zeros((*water_mask.shape, 4))
    overlay[water_mask, 0] = 0.2  # R
    overlay[water_mask, 1] = 0.5  # G
    overlay[water_mask, 2] = 1.0  # B
    overlay[water_mask, 3] = 0.5  # Alpha

    ax.imshow(overlay)

    # 경계선 추가
    from scipy import ndimage
    edges = ndimage.binary_dilation(water_mask, iterations=2) & ~water_mask
    edge_overlay = np.zeros((*water_mask.shape, 4))
    edge_overlay[edges, 0] = 0.0
    edge_overlay[edges, 1] = 1.0
    edge_overlay[edges, 2] = 1.0
    edge_overlay[edges, 3] = 1.0
    ax.imshow(edge_overlay)

    ax.axis('off')
    ax.set_title('Water Body Detection Result', fontsize=14, fontweight='bold')

    # 범례
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor='#3380ff', alpha=0.5, edgecolor='cyan', label='Water Body')
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=10)

    plt.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


@bp.route('/predict', methods=['POST'])
def predict():
    """
    업로드된 이미지에서 Water Body 추출
    POST /api/water_body_segmentation/predict
    Form Data: image (file)
    """
    try:
        # 이미지 파일 확인
        if 'image' not in request.files:
            return jsonify({
                'status': 'error',
                'message': '이미지 파일이 필요합니다.'
            }), 400

        file = request.files['image']
        if file.filename == '':
            return jsonify({
                'status': 'error',
                'message': '파일이 선택되지 않았습니다.'
            }), 400

        # 이미지 로드
        image = Image.open(file.stream)
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # numpy 배열로 변환
        image_array = np.array(image)

        # Water Body 세그멘테이션 (현재는 시뮬레이션)
        # TODO: 실제 DeepLabV3 모델로 교체
        # if deeplabv3_model is not None:
        #     water_mask = run_deeplabv3_prediction(deeplabv3_model, image_array)
        # else:
        water_mask = simulate_water_segmentation(image_array)

        # 결과 이미지 생성
        output_image = generate_prediction_overlay(image_array, water_mask)

        # 통계 계산
        water_pixel_count = np.sum(water_mask)
        total_pixels = water_mask.size
        water_ratio = (water_pixel_count / total_pixels) * 100

        # 연결된 영역 수 계산 (Water Body 개수)
        from scipy import ndimage
        labeled_array, num_features = ndimage.label(water_mask)

        # 면적 추정 (픽셀 기반, 실제 스케일은 이미지 메타데이터 필요)
        # 가정: 1 pixel = 약 1 m² (해상도에 따라 조정 필요)
        pixel_area = 1.0  # m²/pixel
        estimated_area = water_pixel_count * pixel_area

        # 신뢰도 (시뮬레이션)
        confidence = min(95.0, 70.0 + water_ratio * 0.5)

        return jsonify({
            'status': 'success',
            'message': 'Water Body 추출 완료',
            'output_image': output_image,
            'data': {
                'area': f'{estimated_area:,.0f} ㎡',
                'ratio': f'{water_ratio:.1f}%',
                'count': num_features,
                'confidence': f'{confidence:.1f}%',
                'pixels': int(water_pixel_count)
            }
        })

    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500


@bp.route('/load-model', methods=['POST'])
def load_deeplabv3_model():
    """
    DeepLabV3 모델 로드
    POST /api/water_body_segmentation/load-model
    JSON: { "model_path": "path/to/model.h5" }
    """
    global deeplabv3_model

    try:
        data = request.get_json()
        model_path = data.get('model_path')

        if not model_path:
            return jsonify({
                'status': 'error',
                'message': '모델 경로가 필요합니다.'
            }), 400

        # TODO: 실제 모델 로드
        # from app.AI_modules.DeepLabV3 import load_model
        # deeplabv3_model = load_model(model_path)

        return jsonify({
            'status': 'success',
            'message': f'모델 로드 완료: {model_path}'
        })

    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500