model/src/data/hierarchical_embedding.py

"""
계층적 임베딩 처리 관련 함수
"""

import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder
import yfinance as yf

def combine_stocks_for_embedding(all_data, tickers):
    """
    종목 임베딩을 위한 데이터 통합
    """
    print("\n종목 임베딩을 위한 데이터 통합 중...")
    
    all_stocks_data = []
    
    for ticker in tickers:
        if ticker in all_data:
            # 현재 종목 데이터 복사
            stock_df = all_data[ticker].copy()
            
            # 종목 식별자 컬럼 추가
            stock_df['ticker'] = ticker
            
            # 리스트에 추가
            all_stocks_data.append(stock_df)
            
            print(f"{ticker} 데이터: {len(stock_df)}행 처리")
    
    # 모든 데이터프레임을 수직으로 병합
    combined_df = pd.concat(all_stocks_data, axis=0)
    
    combined_df = combined_df.sort_index()
    combined_df = combined_df.reset_index()
    
    # ticker 컬럼을 앞쪽으로 이동
    cols = combined_df.columns.tolist()
    cols.remove('ticker')
    cols = ['ticker'] + cols
    combined_df = combined_df[cols]
    
    return combined_df

def get_industry_data(tickers):
    """
    종목 리스트의 산업 및 섹터 정보를 가져옵니다.
    """
    industry_data = {}
    
    # print("산업 분류 정보 수집 중...")
    for ticker in tickers:
        try:
            # 종목 정보 가져오기
            stock = yf.Ticker(ticker)
            info = stock.info
            
            # 산업 및 섹터 정보 추출
            industry_data[ticker] = {
                'sector': info.get('sector', 'Unknown'),
                'industry': info.get('industry', 'Unknown')
            }
            
            print(f"{ticker}: {industry_data[ticker]['sector']} - {industry_data[ticker]['industry']}")
        except Exception as e:
            print(f"{ticker} 데이터 가져오기 실패: {str(e)}")
            industry_data[ticker] = {'sector': 'Unknown', 'industry': 'Unknown'}
    
    # 데이터프레임으로 변환
    df = pd.DataFrame.from_dict(industry_data, orient='index')
    df.index.name = 'ticker'
    df.reset_index(inplace=True)
    
    return df

def add_industry_encoding(data, industry_data):
    """
    데이터에 산업 정보를 추가하고 인코딩합니다.
    """
    print("\n산업 정보 인코딩 중...")
    
    # 산업 정보 병합
    data = data.merge(industry_data, on='ticker', how='left')
    
    # 누락된 산업 정보는 'Unknown'으로 설정
    data['sector'] = data['sector'].fillna('Unknown')
    data['industry'] = data['industry'].fillna('Unknown')
    
    # 인코딩
    sector_encoder = LabelEncoder()
    industry_encoder = LabelEncoder()
    
    data['sector_id'] = sector_encoder.fit_transform(data['sector'])
    data['industry_id'] = industry_encoder.fit_transform(data['industry'])
    
    # 인코더 정보 반환
    encoders = {
        'sector_encoder': sector_encoder,
        'industry_encoder': industry_encoder,
        'n_sectors': len(sector_encoder.classes_),
        'n_industries': len(industry_encoder.classes_)
    }
    
    print(f"섹터 수: {encoders['n_sectors']}, 산업 수: {encoders['n_industries']}")
    
    return data, encoders

def create_sector_industry_mapping(ticker_list, sector_industry_df):
    """
    종목별 섹터/산업 매핑을 생성합니다.
    """
    sector_mapping = {}
    industry_mapping = {}
    
    for ticker in ticker_list:
        if ticker in sector_industry_df.index:
            sector_mapping[ticker] = int(sector_industry_df.loc[ticker, 'sector_id'])
            industry_mapping[ticker] = int(sector_industry_df.loc[ticker, 'industry_id'])
        else:
            # 정보가 없는 경우 기본값 0 사용
            sector_mapping[ticker] = 0
            industry_mapping[ticker] = 0
    
    return sector_mapping, industry_mapping

def apply_sector_industry_mapping(ticker_data, ticker_encoder, sector_mapping, industry_mapping):
    """
    티커 데이터에 섹터/산업 매핑을 적용합니다.
    """
    # 출력 배열 초기화
    sector_data = np.zeros_like(ticker_data)
    industry_data = np.zeros_like(ticker_data)
    
    # 티커 ID -> 실제 티커 매핑
    id_to_ticker = {v: k for k, v in ticker_encoder.mapping.items()}
    
    # 각 위치에 매핑 적용
    for i in range(ticker_data.shape[0]):
        for j in range(ticker_data.shape[1]):
            ticker_id = ticker_data[i, j]
            if ticker_id in id_to_ticker:
                ticker = id_to_ticker[ticker_id]
                sector_data[i, j] = sector_mapping.get(ticker, 0)
                industry_data[i, j] = industry_mapping.get(ticker, 0)
    
    print(f"적용된 섹터 클래스: {np.unique(sector_data)}")
    print(f"적용된 산업 클래스: {np.unique(industry_data)}")
    return sector_data, industry_data