"""
주식 예측 모델 학습 및 최적화 스크립트
"""
import sys
import os
import argparse
import numpy as np
import json
import pickle
import pandas as pd
from pathlib import Path
from sklearn.preprocessing import LabelEncoder

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.optimization.grid_search import run_optimization_pipeline
from src.optimization.utils import get_project_root

def load_processed_data(tickers):
    """
    전처리된 데이터와 메타 정보를 로드합니다.
    """
    tickers_path = tickers.replace(',', '_') if ',' in tickers else tickers
    data_dir = get_project_root() / "data"
    processed_dir = data_dir / "processed"
    
    processed_path = processed_dir / f"{tickers_path}_processed.pkl"
    encoder_path = processed_dir / f"{tickers_path}_encoder_info.json"
    metadata_path = processed_dir / f"{tickers_path}_metadata.json"
    
    if not processed_path.exists():
        print(f"오류: 전처리된 데이터 파일이 없습니다: {processed_path}")
        return None, None, None
    
    # 데이터 로드
    with open(processed_path, 'rb') as f:
        processed_data = pickle.load(f)
    
    # 원본 데이터 추가
    raw_data_path = data_dir / f"{tickers_path}_data.csv"
    if raw_data_path.exists():
        raw_data = pd.read_csv(raw_data_path)
        processed_data['data'] = raw_data
    
    # 인코더 정보 로드
    encoder_info = {}
    if encoder_path.exists():
        with open(encoder_path, 'r') as f:
            encoder_info = json.load(f)
    
    # 메타데이터 로드
    metadata = {}
    if metadata_path.exists():
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        # 필수 정보만 출력
        print(f"\n데이터셋: {', '.join(metadata.get('tickers', []))}")
        print(f"기간: {metadata.get('start_date', '')} ~ {metadata.get('end_date', '')}")
        print(f"특성 수: {metadata.get('feature_count', '')}, 윈도우: {metadata.get('window_size', '')}")
    
    return processed_data, encoder_info, metadata

def main():
    parser = argparse.ArgumentParser(description="주식 예측 모델 최적화 및 평가")
    parser.add_argument('--tickers', type=str, default='NFLX,TSLA,NVDA,AMD,INTC',
                      help='대상 종목 (콤마로 구분)')
    parser.add_argument('--save', action='store_true', default=True,
                      help='최적 모델 저장 여부')
    parser.add_argument('--metric', type=str, default='combined_score',
                      choices=['combined_score', 'avg_ticker_sharpe', 'sharpe_ratio', 'total_return'],
                      help='최적화 기준 지표')
    parser.add_argument('--output', type=str, default='grid_search_results.pkl',
                      help='결과 파일 경로')
    parser.add_argument('--model_output', type=str, default='best_contime',
                      help='TensorFlow Lite 모델 저장 경로')
    parser.add_argument('--visualize', action='store_true', default=True,
                      help='시각화 실행 여부')
    
    args = parser.parse_args()
    
    # 전처리된 데이터 로드
    processed_data, encoder_info, metadata = load_processed_data(args.tickers)
    
    if processed_data is None:
        print("전처리된 데이터를 로드할 수 없습니다.")
        return
    
    # 티커 인코더 추출
    ticker_encoder = LabelEncoder()
    if 'ticker_encoder' in encoder_info:
        # JSON 파일에서 로드한 딕셔너리를 LabelEncoder로 변환
        ticker_mapping = encoder_info['ticker_encoder']
        ticker_list = [ticker for _, ticker in sorted([(int(i), ticker) for i, ticker in ticker_mapping.items()])]
        ticker_encoder.fit(ticker_list)
    else:
        # 티커 리스트 추출
        ticker_list = metadata.get('tickers', args.tickers.split(','))
        ticker_encoder.fit(ticker_list)
    
    # 임의의 시드 설정
    np.random.seed(42)
    
    # 최적화 실행
    print("\n모델 최적화 시작...")
    print(f"최적화 기준: {args.metric}")
    
    # 결과 파일 저장 경로 설정
    results_dir = Path(get_project_root()) / "models" / "results"
    results_dir.mkdir(parents=True, exist_ok=True)

    output_path = results_dir / args.output
    model_output = results_dir / args.model_output

    optimization_results = run_optimization_pipeline(
        data_dict=processed_data,
        ticker_encoder=ticker_encoder,
        metric=args.metric,
        output_path=output_path,
        save=args.save,
        model_output=model_output,
        sector_industry_df=processed_data.get('sector_industry_df'),
        run_visualizations=args.visualize
    )
    
    print("\n최적화 완료!")
    
    # 결과 요약
    best_config = optimization_results.get('best_config', {})
    if best_config:
        print("\n최적 설정 요약:")
        for key, value in best_config.items():
            print(f"  {key}: {value}")
    
    return optimization_results

if __name__ == "__main__":
    main()