import numpy as np
import matplotlib.pyplot as plt
import cv2
from io import BytesIO
from PIL import Image
import streamlit as st


try:
    import tensorflow as tf
    from tensorflow.keras.datasets import fashion_mnist

    HAS_TF = True
except ImportError:
    HAS_TF = False

import os
from datetime import datetime


def map_keyword_to_indices(text, class_names):
    """私有方法：从推荐文本解析服装类别索引"""
    keyword_map = {
        # 中文关键词映射
        "t恤": 0, "衬衫": 6, "外套": 4, "套头衫": 2,
        "裤子": 1, "牛仔裤": 1,
        "连衣裙": 3, "裙子": 3,
        "凉鞋": 5, "运动鞋": 7, "短靴": 9, "靴子": 9
    }

    matched_indices = []
    text_lower = text.lower()
    for keyword, idx in keyword_map.items():
        if keyword in text_lower:
            matched_indices.append(idx)
    return list(set(matched_indices))  # 去重


# outfits.py 修改后的代码
def generate_sample_outfits(recommendation_text, save_to_local=False, output_dir="C:/Users/Administrator/Desktop"):
    """生成示例穿搭图片（包含连衣裙选项）并可选保存到本地"""
    outfits = []

    # 创建输出目录
    if save_to_local and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if HAS_TF:
        try:
            (train_images, train_labels), _ = fashion_mnist.load_data()
            class_names = ['T恤', '裤子', '套头衫', '连衣裙', '外套',
                           '凉鞋', '衬衫', '运动鞋', '包', '短靴']

            # 从推荐文本中提取服装类别索引
            matched_indices = map_keyword_to_indices(recommendation_text, class_names)

            # 根据匹配结果生成图片
            if not matched_indices:
                st.warning("无法从推荐文本中识别有效服装类别")
                return generate_fallback_outfits(save_to_local, output_dir)

            # 生成3套穿搭方案
            for i in range(3):
                try:
                    # 随机决定是生成连衣裙组合还是上下装组合
                    use_dress = np.random.choice([True, False]) and (3 in matched_indices)

                    if use_dress:
                        # 连衣裙+鞋子组合
                        shoes_indices = [i for i in matched_indices if i in [5, 7, 9]]
                        if not shoes_indices:
                            continue

                        shoes_idx = np.random.choice(shoes_indices)
                        dress_img = train_images[np.random.choice(np.where(train_labels == 3)[0])]
                        shoes_img = train_images[np.random.choice(np.where(train_labels == shoes_idx)[0])]

                        dress_img = cv2.resize(dress_img, (128, 256), interpolation=cv2.INTER_NEAREST)
                        shoes_img = cv2.resize(shoes_img, (128, 128), interpolation=cv2.INTER_NEAREST)
                        combined = np.vstack([dress_img, shoes_img])
                        combined_rgb = np.stack([combined] * 3, axis=-1)

                        title = f"{class_names[3]}+{class_names[shoes_idx]}"
                    else:
                        # 上下装+鞋子组合
                        top_indices = [i for i in matched_indices if i in [0, 2, 4, 6]]
                        bottom_indices = [i for i in matched_indices if i in [1]]
                        shoes_indices = [i for i in matched_indices if i in [5, 7, 9]]

                        if not top_indices or not bottom_indices or not shoes_indices:
                            continue

                        top_idx = np.random.choice(top_indices)
                        bottom_idx = np.random.choice(bottom_indices)
                        shoes_idx = np.random.choice(shoes_indices)

                        top_img = train_images[np.random.choice(np.where(train_labels == top_idx)[0])]
                        bottom_img = train_images[np.random.choice(np.where(train_labels == bottom_idx)[0])]
                        shoes_img = train_images[np.random.choice(np.where(train_labels == shoes_idx)[0])]

                        top_img = cv2.resize(top_img, (128, 128), interpolation=cv2.INTER_NEAREST)
                        bottom_img = cv2.resize(bottom_img, (128, 128), interpolation=cv2.INTER_NEAREST)
                        shoes_img = cv2.resize(shoes_img, (128, 128), interpolation=cv2.INTER_NEAREST)
                        combined = np.vstack([top_img, bottom_img, shoes_img])
                        combined_rgb = np.stack([combined] * 3, axis=-1)

                        title = f"{class_names[top_idx]}+{class_names[bottom_idx]}+{class_names[shoes_idx]}"

                    # 直接使用PIL保存图片，避免plt保存问题
                    img = Image.fromarray(combined_rgb)

                    if save_to_local:
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        filename = f"{output_dir}/outfit_{timestamp}_{i}.png"
                        img.save(filename, format='PNG', quality=95)

                    outfits.append(img)

                except Exception as e:
                    st.warning(f"生成第{i + 1}套穿搭图片时出错: {str(e)}")
                    outfits.append(generate_single_fallback_outfit(i + 1, save_to_local, output_dir))
                    continue

            return outfits if outfits else generate_fallback_outfits(save_to_local, output_dir)

        except Exception as e:
            st.warning(f"生成服装图片出错: {str(e)}，使用文字替代")
            return generate_fallback_outfits(save_to_local, output_dir)

    return generate_fallback_outfits(save_to_local, output_dir)


def generate_single_fallback_outfit(index, save_to_local=False, output_dir=None):
    """生成单个回退的文字穿搭方案"""
    fig, ax = plt.subplots(figsize=(3, 4))
    ax.set_facecolor('#f0f2f6')
    ax.text(0.5, 0.5, f"穿搭方案 {index}\n(图片生成失败)",
            ha='center', va='center', fontsize=10)
    ax.axis('off')

    buf = BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1, dpi=100)
    buf.seek(0)
    img = Image.open(buf)
    plt.close(fig)

    if save_to_local and output_dir:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{output_dir}/outfit_{timestamp}_{index}_fallback.png"
        img.save(filename)

    return img


def generate_fallback_outfits(save_to_local=False, output_dir="C:/Users/Administrator/Desktop"):
    """生成回退的文字穿搭方案"""
    outfits = []
    for i in range(3):
        outfits.append(generate_single_fallback_outfit(i + 1, save_to_local, output_dir))
    return outfits