# -*- coding: utf-8 -*-
import os as _o
# chdir กลับ project root — เฉพาะตอนที่อยู่ในโครงโปรเจกต์จริง (บน Docker ไฟล์ถูกวางแบนๆ ต้องไม่ chdir)
_root = _o.path.dirname(_o.path.dirname(_o.path.dirname(_o.path.abspath(__file__))))
if _o.path.isdir(_o.path.join(_root, 'results')):
    _o.chdir(_root)
# ============================================================
# features.py — รูป 1 ใบ -> 22 ตัวเลข (7 ratio + 15 deep)
#
# ===== ต้องการอะไร =====
#   นี่คือ "ท่อต้นน้ำ" ที่ยังขาดอยู่ชิ้นเดียว
#   predict_bmi.py กับ pipeline.py ทำงานได้แล้ว แต่ต้องป้อน 22 ตัวเลขให้มันเอง
#   ไฟล์นี้ทำให้ป้อน "รูป" ได้เลย = แอปรับรูปจากผู้ใช้ได้จริง
#
# ===== ไฟล์นี้ทำอะไร =====
#   รับ path รูป -> คืน (ratio7, deep15) พร้อมส่งต่อให้ predict_bmi
#   ทำ 2 สายพร้อมกัน:
#     สาย A: detectron2 (หาคน+ข้อต่อ) -> Pose2Seg (แยกร่างจากพื้นหลัง) -> วัดสัดส่วน -> 7 ratio
#     สาย B: DenseNet121 ที่ fine-tune แล้ว -> 15 deep features
#
#   *** ไฟล์นี้รันบน Windows ไม่ได้ ***
#   detectron2 ต้อง Linux -> ใช้งานจริงบน Hugging Face Spaces (Docker) เท่านั้น
#   ที่เหลือ (nutrition/predict_bmi/pipeline) รันบน Windows ได้หมด จึงแยกไฟล์นี้ออกมา
#
# ===== วิธีคิด (สำคัญที่สุดในไฟล์นี้) =====
#   ทุกขั้นตอนต้องทำ "เหมือนตอนเทรนบน Kaggle เป๊ะ" ไม่งั้นตัวเลขที่ได้จะเป็นคนละชุด
#   กับที่ SVR เคยเห็น -> ทำนาย BMI มั่วโดยไม่มี error ฟ้อง (พังเงียบ)
#   จุดที่ต้องเหมือนเป๊ะ:
#     1. เตรียมรูป: ย่อด้านยาว=224 + เติมขอบ "ดำ" + normalize ค่า ImageNet
#     2. โมเดลหาข้อต่อ: keypoint_rcnn_R_101_FPN_3x  (R_101 ไม่ใช่ R_50)
#     3. Detected.py ต้องเป็นตัวที่ patch แล้ว (ตัด SCHP/PSP/CPM/CRFRNN)
#        -> mask จะ "มีแขนติดอยู่" ซึ่งตอนเทรนก็มีแขน = ต้องปล่อยให้มีเหมือนกัน
# ============================================================
import os
import numpy as np

# ---------- ที่อยู่ไฟล์ (ตั้งผ่าน environment variable ได้ เพื่อให้ Docker เปลี่ยนง่าย) ----------
MODEL_DIR    = os.environ.get('MODEL_DIR', 'results/dense')
DENSE_PT     = os.environ.get('DENSE_PT', os.path.join(MODEL_DIR, 'dense_bmi_best.pt'))
POSE2SEG_PKL = os.environ.get('POSE2SEG_PKL', 'models/pose2seg_release.pkl')
BMI_REPO     = os.environ.get('BMI_REPO', '2DImage2BMI-main')   # โฟลเดอร์ที่มี Detected.py (patch แล้ว)

# ===== แปลงทุก path เป็น absolute ทันทีตอน import =====
# ทำไมต้องทำ: _load_anthro() ต้อง os.chdir() เข้าไปใน repo (เพราะ Detected.py อ้าง path แบบ relative ข้างใน)
# ถ้าปล่อย path เป็น relative ไว้ พอ chdir แล้ว path เดิมจะชี้ผิดที่ทันที -> หาไฟล์ไม่เจอแบบงงๆ
DENSE_PT     = os.path.abspath(DENSE_PT)
POSE2SEG_PKL = os.path.abspath(POSE2SEG_PKL)
BMI_REPO     = os.path.abspath(BMI_REPO)

# ---------- ชื่อ+ลำดับคอลัมน์ (ต้องตรงกับ bmi_bundle.pkl เป๊ะ ห้ามสลับ) ----------
RATIO_COLS = ['WSR', 'WTR', 'WHpR', 'WHdR', 'HpHdR', 'Area', 'H2W']
N_DEEP = 15

# ---------- ค่า normalize ของ ImageNet (ตัวเดียวกับตอนเทรน) ----------
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]
IMG_SIZE = 224

_net = None        # DenseNet — โหลดครั้งเดียวแล้วใช้ซ้ำ
_proc = None       # detectron2 + Pose2Seg — ตัวนี้หนักมาก ยิ่งต้องโหลดครั้งเดียว


# ============================================================
# สาย B (1/2): เตรียมรูปให้เหมือนตอนเทรน
# ============================================================
def _resize_pad(pil_img, size=IMG_SIZE):
    """
    ย่อรูปให้ "ด้านที่ยาวที่สุด" = 224 แล้วเติมขอบดำจนเป็นสี่เหลี่ยม 224x224

    ทำไมไม่ resize ตรงๆ เป็น 224x224:
      รูปคนส่วนใหญ่เป็นแนวตั้ง (สูงมากกว่ากว้าง) ถ้าบีบเป็นจัตุรัสตรงๆ
      คนผอมจะกลายเป็นคนอ้วนในรูป -> CNN อ่านสัดส่วนผิดหมด
      การเติมขอบดำรักษาสัดส่วนตัวจริงไว้ (เหมือนดูหนังจอกว้างมีแถบดำบนล่าง)
    """
    from PIL import Image
    w, h = pil_img.size
    k = size / max(w, h)                                          # อัตราย่อ ยึดด้านที่ยาวสุด
    img = pil_img.resize((max(1, int(w * k)), max(1, int(h * k))))
    canvas = Image.new('RGB', (size, size), (0, 0, 0))            # ผ้าใบดำ 224x224
    canvas.paste(img, ((size - img.width) // 2, (size - img.height) // 2))   # วางไว้ตรงกลาง
    return canvas


def _get_transform():
    """สร้างชุดแปลงรูป: resize+pad -> tensor -> normalize (เหมือนตอนเทรนทุกขั้น)"""
    from torchvision import transforms
    return transforms.Compose([
        transforms.Lambda(_resize_pad),
        transforms.ToTensor(),                                    # 0-255 -> 0.0-1.0
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),        # ปรับให้ตรงกับที่ DenseNet เคยเรียนมา
    ])


# ============================================================
# สาย B (2/2): โครงโมเดล DenseNet — ต้องเหมือนตอนเทรน ไม่งั้นโหลด .pt ไม่ได้
# ============================================================
def _build_densenet():
    """
    สร้างโครงเปล่าให้เหมือนตอนเทรน แล้วค่อยเทค่าน้ำหนักจาก .pt ใส่

    โครง: DenseNet121 -> เปลี่ยนชั้นท้ายให้ออก 15 ตัว -> head(15->1) ทำนาย BMI
    ตอนใช้งานจริงเราเอาแค่ "15 ตัวกลาง" ทิ้ง head
    (head มีไว้ตอนเทรนเพื่อบังคับให้ 15 ตัวนั้นมีความหมายเกี่ยวกับ BMI)
    """
    import torch.nn as nn
    from torchvision import models

    class DenseBMI(nn.Module):
        def __init__(self):
            super().__init__()
            self.backbone = models.densenet121(weights=None)      # weights=None เพราะเดี๋ยวโหลดจาก .pt ทับ
            self.backbone.classifier = nn.Linear(self.backbone.classifier.in_features, N_DEEP)
            self.head = nn.Linear(N_DEEP, 1)

        def forward(self, x):
            f = self.backbone(x)          # -> 15 deep features
            return self.head(f), f        # คืน (BMI ที่ทำนาย, 15 features)

    return DenseBMI()


def _load_densenet():
    """โหลด DenseNet ครั้งแรกครั้งเดียว (lazy load) — ครั้งต่อไปใช้ตัวเดิม"""
    global _net
    if _net is not None:
        return _net

    import torch
    if not os.path.exists(DENSE_PT):
        raise FileNotFoundError(f'ไม่เจอไฟล์โมเดล {DENSE_PT} — ตั้ง env DENSE_PT ให้ถูก หรือก็อปไฟล์มาวาง')

    net = _build_densenet()
    state = torch.load(DENSE_PT, map_location='cpu')              # map_location='cpu' = ใช้ได้แม้ไม่มี GPU
    net.load_state_dict(state)
    net.eval()                                                    # โหมดใช้งาน (ปิด dropout/batchnorm update)
    _net = net
    return _net


# ============================================================
# สาย A: detectron2 + Pose2Seg -> 7 ratio
# ============================================================
def _load_anthro():
    """
    โหลดตัวสกัดสัดส่วน (หนักสุดในระบบ — โหลดครั้งเดียวพอ)

    ⚠️ ตัวนี้ต้อง Linux เท่านั้น (detectron2 build บน Windows ไม่ผ่าน)
    """
    global _proc
    if _proc is not None:
        return _proc

    import sys
    # ต้อง chdir เข้า repo ก่อน เพราะ Detected.py อ้าง path แบบ relative ข้างใน
    if not os.path.isdir(BMI_REPO):
        raise FileNotFoundError(f'ไม่เจอโฟลเดอร์ {BMI_REPO} (ที่มี Detected.py) — ตั้ง env BMI_REPO ให้ถูก')
    sys.path.insert(0, os.path.abspath(BMI_REPO))
    os.chdir(BMI_REPO)

    try:
        from Detected import Image_Processor
    except ImportError as e:
        raise ImportError(
            f'import Detected ไม่ผ่าน: {e}\n'
            '  เช็ค 3 อย่าง:\n'
            '   1) Detected.py เป็นตัวที่ patch แล้วหรือยัง (ต้องตัด import HumanParser/PSP/CPM/CRFRNN)\n'
            '   2) เอา lib/ ของ Pose2Seg มาเติมใน repo แล้วหรือยัง (lib.transforms, lib.timer)\n'
            '   3) numpy เป็น 1.26.4 หรือยัง (numpy 2.0 ลบ np.mat -> Pose2Seg พัง)'
        ) from e

    if not os.path.exists(POSE2SEG_PKL):
        raise FileNotFoundError(f'ไม่เจอ {POSE2SEG_PKL} (โมเดล Pose2Seg)')

    # ===== หัวใจ: R_101 ไม่ใช่ R_50 — ต้องตรงกับตอนเทรน ไม่งั้นจุดข้อต่อเลื่อน ratio เพี้ยน =====
    _proc = Image_Processor(POSE2SEG_PKL, 'COCO-Keypoints/keypoint_rcnn_R_101_FPN_3x.yaml')
    return _proc


def extract_ratio(bgr_image):
    """
    รับ  : รูปแบบ BGR (numpy array จาก cv2.imread)
    คืน  : list 7 ค่า เรียงตาม RATIO_COLS

    ข้างใน Process() จะ: หาคน -> หาข้อต่อ 17 จุด -> วาด mask -> วัดความกว้างที่ระดับต่างๆ
    แล้วคำนวณเป็นอัตราส่วน (อัตราส่วนไม่ขึ้นกับระยะถ่ายใกล้/ไกล จึงใช้แทนหน่วย cm ได้)
    """
    P = _load_anthro()
    try:
        F = P.Process(bgr_image)
    except Exception as e:
        # เคสที่เจอบ่อยสุด: รูปไม่มีคน / คนถูกบัง / มีหลายคนจนสับสน
        raise RuntimeError(
            f'สกัดสัดส่วนจากรูปไม่สำเร็จ ({type(e).__name__}: {e})\n'
            '  มักเกิดจาก: ในรูปไม่มีคน / เห็นไม่เต็มตัว / มีหลายคน / รูปมืดเกินไป\n'
            '  แนะนำผู้ใช้: ถ่ายเต็มตัว ยืนตรง พื้นหลังโล่ง คนเดียวในรูป'
        ) from e

    return [float(getattr(F, c)) for c in RATIO_COLS]   # ดึงตามลำดับ RATIO_COLS เป๊ะ


def extract_deep(pil_image):
    """
    รับ  : รูป PIL (RGB)
    คืน  : list 15 ค่า (deep features จาก DenseNet ที่ fine-tune แล้ว)
    """
    import torch
    net = _load_densenet()
    tf = _get_transform()

    x = tf(pil_image).unsqueeze(0)                # เพิ่มมิติ batch -> (1, 3, 224, 224)
    with torch.no_grad():                         # ไม่ต้องคิดอนุพันธ์ = เร็วขึ้น ประหยัดแรม
        _, feat = net(x)                          # เอาแค่ feat ทิ้งค่าทำนาย BMI ของ head
    return feat.cpu().numpy()[0].astype(float).tolist()


# ============================================================
# หัวใจ: รูป -> 22 ตัวเลข (ฟังก์ชันเดียวที่ข้างนอกเรียก)
# ============================================================
def image_to_features(image_path):
    """
    รับ  : path รูป
    คืน  : dict {'ratio7': [...7], 'deep15': [...15]}

    เอาไปต่อกับ pipeline.analyze(ratio7, deep15, ...) ได้ทันที
    """
    import cv2
    from PIL import Image

    if not os.path.exists(image_path):
        raise FileNotFoundError(f'ไม่เจอรูป: {image_path}')

    # ---------- สาย A: ต้องใช้ BGR (cv2) ----------
    bgr = cv2.imread(image_path)
    if bgr is None:
        raise ValueError(f'อ่านรูปไม่ได้ (ไฟล์เสียหรือไม่ใช่รูป?): {image_path}')
    ratio7 = extract_ratio(bgr)

    # ---------- สาย B: ต้องใช้ RGB (PIL) ----------
    # ⚠️ ระวัง: cv2 อ่านมาเป็น BGR ส่วน PIL เป็น RGB — สลับกันเมื่อไหร่สีเพี้ยน CNN อ่านผิด
    pil = Image.open(image_path).convert('RGB')
    deep15 = extract_deep(pil)

    # ---------- ตรวจว่าได้ครบจำนวนจริง ----------
    if len(ratio7) != len(RATIO_COLS):
        raise RuntimeError(f'ratio ควรได้ {len(RATIO_COLS)} ค่า แต่ได้ {len(ratio7)}')
    if len(deep15) != N_DEEP:
        raise RuntimeError(f'deep ควรได้ {N_DEEP} ค่า แต่ได้ {len(deep15)}')

    # ---------- ดักค่าเสีย (inf/nan) ที่จะทำให้ SVR พังหรือทายมั่ว ----------
    allv = np.array(ratio7 + deep15, dtype=float)
    if not np.isfinite(allv).all():
        bad = [RATIO_COLS[i] for i in range(len(ratio7)) if not np.isfinite(ratio7[i])]
        raise RuntimeError(f'ได้ค่าเสีย (inf/nan) จากรูปนี้ ตัวที่เสีย: {bad or "deep features"}\n'
                           '  มักเกิดตอนวัดความกว้างได้ 0 แล้วเอาไปหาร -> ถ่ายใหม่ให้เห็นเต็มตัว')

    return {'ratio7': ratio7, 'deep15': deep15}


# ============================================================
# เทสตัวเอง
#   บน Windows: จะเช็คได้แค่ "สาย B (DenseNet)" เท่านั้น
#   สาย A ต้องรอ Linux — จะขึ้นข้อความบอกชัดว่าเช็คไม่ได้ ไม่ใช่พัง
#
#   วิธีใช้: python scripts/app/features.py <path รูป>
# ============================================================
if __name__ == '__main__':
    import sys

    print('=' * 66)
    print('  เทส features.py')
    print('=' * 66)

    # ---------- เช็คว่าของที่ต้องใช้มีครบไหม ----------
    print('\n  [1] เช็คไฟล์ที่ต้องใช้:')
    for label, p in [('DenseNet .pt', DENSE_PT), ('Pose2Seg .pkl', POSE2SEG_PKL), ('repo Detected.py', BMI_REPO)]:
        print(f'      {"✅" if os.path.exists(p) else "❌"} {label:18s} {p}')

    # ---------- เช็คว่า library ลงครบไหม ----------
    print('\n  [2] เช็ค library:')
    for mod, note in [('torch', 'สาย B'), ('torchvision', 'สาย B'), ('cv2', 'อ่านรูป'),
                      ('detectron2', 'สาย A — Linux เท่านั้น')]:
        try:
            __import__(mod)
            print(f'      ✅ {mod:14s} ({note})')
        except ImportError:
            print(f'      ❌ {mod:14s} ({note}) — ยังไม่ได้ลง')

    # ---------- ถ้าใส่ path รูปมา ให้ลองสกัดจริง ----------
    if len(sys.argv) > 1:
        img = sys.argv[1]
        print(f'\n  [3] ลองสกัดจากรูปจริง: {img}')
        try:
            f = image_to_features(img)
            print(f'      ✅ ratio7  = {[round(v, 4) for v in f["ratio7"]]}')
            print(f'      ✅ deep15  = {[round(v, 3) for v in f["deep15"][:5]]} ... (อีก 10 ตัว)')
            print('\n      ต่อเข้า pipeline ได้เลย:')
            print('        from pipeline import analyze')
            print('        analyze(f["ratio7"], f["deep15"], height_cm=170, age=25, sex=1)')
        except Exception as e:
            print(f'      ❌ {type(e).__name__}: {e}')
    else:
        print('\n  [3] ไม่ได้ใส่ path รูป — ข้ามการสกัดจริง')
        print('      วิธีลอง: python scripts/app/features.py <path รูป>')

    print('\n' + '=' * 66)
    print('  หมายเหตุ: ถ้า detectron2 ขึ้น ❌ = เครื่องนี้รันสาย A ไม่ได้ (ปกติสำหรับ Windows)')
    print('           ต้องรันบน Linux/Docker เท่านั้น')
    print('=' * 66)