"""
주식 데이터 수집 및 전처리 스크립트
"""
import sys
import os
import argparse
import numpy as np
import json
import pickle
from datetime import datetime
from pathlib import Path

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.data.data_integration import process_stock_data
from src.data.processors import prepare_data

def main():
    parser = argparse.ArgumentParser(description="주식 데이터 수집 및 전처리")
    parser.add_argument('--tickers', type=str, default='NFLX,TSLA,NVDA,AMD,INTC',
                      help='수집할 종목 (쉼표로 구분)')
    parser.add_argument('--start_date', type=str, default='2020-01-01',
                      help='시작 날짜 (YYYY-MM-DD)')
    parser.add_argument('--end_date', type=str, default=None,
                      help='종료 날짜 (YYYY-MM-DD), 미지정시 오늘')
    parser.add_argument('--fred_api_key', type=str, 
                      default='4c55d0ee6170369793707da4cba1b7be',
                      help='FRED API 키')
    parser.add_argument('--window_size', type=int, default=60,
                      help='시계열 윈도우 크기')
    
    args = parser.parse_args()
    
    # 종목 리스트 파싱
    training_tickers = args.tickers.split(',')
    
    # 날짜 설정
    start_date = args.start_date
    end_date = args.end_date or datetime.today().strftime('%Y-%m-%d')
    
    print(f"데이터 수집 시작: {', '.join(training_tickers)}")
    print(f"기간: {start_date} ~ {end_date}")
    
    # 데이터 처리 및 통합
    final_data, all_data, industry_encoders = process_stock_data(
        training_tickers, 
        start_date, 
        end_date, 
        args.fred_api_key
    )
    
    # 데이터셋 요약 정보 출력
    print("\n===== 통합 데이터셋 요약 =====")
    print(f"크기: {final_data.shape[0]}행, {final_data.shape[1]}열")
    if 'ticker' in final_data.columns:
        print(f"종목 수: {final_data['ticker'].nunique()}개")
    if 'sector' in final_data.columns:
        print(f"섹터 수: {final_data['sector'].nunique()}개")
    if 'industry' in final_data.columns:
        print(f"산업 수: {final_data['industry'].nunique()}개")
    print(f"기간: {final_data['Date'].min()} ~ {final_data['Date'].max()}")
    
    # 종목별 정보 출력
    print(f"종목별 행 수:")
    for ticker in training_tickers:
        if 'ticker' in final_data.columns:
            count = len(final_data[final_data['ticker'] == ticker])
        else:
            count = "정보 없음"
        print(f"  - {ticker}: {count}행")
    
    # 저장 디렉토리 준비
    os.makedirs("./data", exist_ok=True)
    
    data_dir = Path("./data")
    all_tickers = '_'.join(training_tickers)
    raw_filename = data_dir / f"{all_tickers}_data.csv"
    
    # 원본 데이터 저장
    final_data.to_csv(raw_filename, index=False)
    print(f"원본 데이터가 저장되었습니다: {raw_filename}")
    
    # 데이터 전처리 및 모델 학습용 데이터 생성
    if 'Date' in final_data.columns:
        final_data = final_data.set_index('Date')
        final_data.sort_index(inplace=True)
    
    # 데이터 전처리 실행
    print("\n전처리 및 학습용 데이터 생성 시작...")
    processed_data, ticker_encoder, _ = prepare_data(final_data, window_size=args.window_size)
    
    # 섹터 및 산업 매핑 (있는 경우)
    sector_map = {}
    industry_map = {}
    if 'sector' in final_data.columns and 'sector_id' in final_data.columns:
        for sector, sector_id in zip(final_data['sector'].unique(), final_data['sector_id'].unique()):
            sector_map[sector] = int(sector_id)
    
    if 'industry' in final_data.columns and 'industry_id' in final_data.columns:
        for industry, industry_id in zip(final_data['industry'].unique(), final_data['industry_id'].unique()):
            industry_map[industry] = int(industry_id)
    
    # 결과에 섹터 매핑 추가
    processed_data['sector_map'] = sector_map
    processed_data['industry_map'] = industry_map
    processed_data['feature_count'] = processed_data['x_train'].shape[2]
    
    # 전처리된 데이터 저장 (pickle 형식)
    processed_dir = data_dir / "processed"
    processed_dir.mkdir(parents=True, exist_ok=True)
    
    processed_filename = processed_dir / f"{all_tickers}_processed.pkl"
    with open(processed_filename, 'wb') as f:
        pickle.dump(processed_data, f)
    print(f"전처리된 데이터가 저장되었습니다: {processed_filename}")
    
    # 인코더 정보 저장 (JSON 형식)
    encoder_info = {
        'ticker_encoder': {str(i): ticker for i, ticker in enumerate(ticker_encoder.classes_)},
        'sector_map': processed_data.get('sector_map', {}),
        'industry_map': processed_data.get('industry_map', {})
    }
    
    encoder_filename = processed_dir / f"{all_tickers}_encoder_info.json"
    with open(encoder_filename, 'w') as f:
        json.dump(encoder_info, f)
    print(f"인코더 정보가 저장되었습니다: {encoder_filename}")
    
    # 메타데이터 저장
    metadata = {
        'feature_count': processed_data.get('feature_count', 0),
        'window_size': args.window_size,
        'tickers': training_tickers,
        'start_date': str(processed_data.get('start_date', start_date)),
        'end_date': str(processed_data.get('end_date', end_date)),
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    
    metadata_filename = processed_dir / f"{all_tickers}_metadata.json"
    with open(metadata_filename, 'w') as f:
        json.dump(metadata, f)
    print(f"메타데이터가 저장되었습니다: {metadata_filename}")
    
    print("\n===== 전처리 완료 =====")
    print(f"학습 데이터 크기: {processed_data['x_train'].shape}")
    print(f"검증 데이터 크기: {processed_data['x_val'].shape}")
    print(f"테스트 데이터 크기: {processed_data['x_test'].shape}")
    print(f"특성 수: {processed_data['feature_count']}")

if __name__ == "__main__":
    main()