# 모듈 임포트
import pandas as pd
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer
from tqdm.notebook import tqdm
from transformers import AutoTokenizer, AutoModel

data_path = "/emotion_korTran.data"

data = []
with open(data_path, "r", encoding="utf-8") as f:
    next(f)  # 첫 줄 건너뜀 (헤더)
    for line in f:
        line = line.strip()
        if not line:
            continue
        id_text = line.rsplit(",", 1)  # 오른쪽에서부터 마지막 콤마 기준으로 분할
        if len(id_text) == 2:
            idx_text, label = id_text
            idx, text = idx_text.split(",", 1)
            data.append((int(idx), text, label))

df = pd.DataFrame(data, columns=["id", "text", "label"])

label_map = {
    'joy': 0,
    'sadness': 1,
    'anger': 2,
    'fear': 3,
    'love': 4,
    'surprise': 5
}

df['label'] = df['label'].map(label_map)

tokenizer = AutoTokenizer.from_pretrained("monologg/kobert")
base_model = AutoModel.from_pretrained("monologg/kobert")

class EmotionDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=64):
        self.texts = texts
        self.labels = labels # 감정 분류를 위한 감정 라벨
        self.tokenizer = tokenizer # BERT 모델 입력을 위해 토큰화 수행
        self.max_len = max_len # 입력 시퀀스를 고정 길이로 맞추기

    def __len__(self):
        return len(self.texts)

    # 모델이 학습 시 한 문장씩 꺼낼 수 있게 함
    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            padding='max_length',
            truncation=True,
            max_length=self.max_len,
            return_tensors="pt"
        )
        item = {key: val.squeeze(0) for key, val in encoding.items()}
        item["labels"] = torch.tensor(self.labels[idx])
        return item

from sklearn.model_selection import train_test_split

# train과 validation을 9:1로 나눔
train_texts, val_texts, train_labels, val_labels = train_test_split(
    df["text"], df["label"], test_size=0.1, random_state=42
)

train_dataset = EmotionDataset(train_texts.tolist(), train_labels.tolist(), tokenizer)
# 학습 데이터 일부만 추출 (10%만) <- Colab 런타인 시간 내 학습이 완료되어야 함
data_size = len(train_dataset)
sampled_dataset = torch.utils.data.Subset(train_dataset, range(data_size//10, data_size // 10 * 2))

val_dataset = EmotionDataset(val_texts.tolist(), val_labels.tolist(), tokenizer)

train_loader = DataLoader(sampled_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32)

# KoBERT Customizing -> PyTorch Model
class KoBERTClassifier(nn.Module):
    def __init__(self, bert, num_classes=6):
        super(KoBERTClassifier, self).__init__()
        self.bert = bert # 사전학습된 KoBERT
        self.classifier = nn.Linear(bert.config.hidden_size, num_classes) # BERT 출력 -> 감정 분류 클래스 수로 줄이는 선형 layer

    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        outputs = self.bert( # KoBERT에 입력 tensor를 넣음
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )
        pooled = outputs.pooler_output
        return self.classifier(pooled)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # GPU 있으면 GPU 사용함.

model = KoBERTClassifier(base_model).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
loss_fn = nn.CrossEntropyLoss()

# 저장된 모델 가중치 불러오기
model.load_state_dict(torch.load("/kobert_emotion_epoch4.pt", map_location=device))  # 파일명 확인
print("모델 가중치 불러오기 완료")

# Optimizer, Loss 재설정
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
loss_fn = nn.CrossEntropyLoss()

# 이어서 학습
EPOCHS = 1  # 추가로 학습할 에폭 수
start_epoch = 8  # 이어서 시작할 에폭 번호

for epoch in range(start_epoch, start_epoch + EPOCHS):
    model.train()
    total_loss = 0

    for batch in tqdm(train_loader):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        token_type_ids = batch['token_type_ids'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask, token_type_ids)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(train_loader)
    print(f"[Epoch {epoch}] Loss: {avg_loss:.4f}")

from google.colab import files

# 에폭별 모델 저장
model_path = f"/kobert_emotion_epoch{epoch+1}.pt"
drive_path = f"/{model_path}"

# 로컬에 저장
torch.save(model.state_dict(), model_path)

# model 다운로드
files.download(model_path)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 모델 구조 다시 정의 (기존과 동일하게)
model = KoBERTClassifier(base_model).to(device)

# 저장된 가중치 불러오기
model_path = "/content/drive/MyDrive/Colab Notebooks/NLP/models/kobert_emotion_epoch8.pt"
model.load_state_dict(torch.load(model_path, map_location=device))

# 평가 모드로 전환
model.eval()

from sklearn.metrics import accuracy_score

def evaluate_accuracy(model, dataloader, tokenizer, sample_texts=None):
    model.eval()
    preds = []
    true_labels = []
    sample_outputs = []

    with torch.no_grad():
        for batch in tqdm(dataloader):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            token_type_ids = batch['token_type_ids'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask, token_type_ids)
            predictions = torch.argmax(outputs, dim=1)

            preds.extend(predictions.cpu().numpy())
            true_labels.extend(labels.cpu().numpy())

            # 샘플 출력용 텍스트 복원
            if sample_texts and len(sample_outputs) < 10:
                texts = sample_texts[len(sample_outputs):len(sample_outputs)+len(predictions)]
                for text, pred, true in zip(texts, predictions, labels):
                    sample_outputs.append((text, pred.item(), true.item()))
                    if len(sample_outputs) == 10:
                        break

    acc = accuracy_score(true_labels, preds)
    print(f"Validation Accuracy: {acc * 100:.2f}%\n")

evaluate_accuracy(model, val_loader, tokenizer)

# 해리포터 소설 불러오기
with open("해리포터_통합본_전처리.txt", "r", encoding="utf-8") as f:
    texts = f.read().split('\n')

# 해리포터 전권을 챕터별로 구분
chapters = []
chapter_titles = []
chapter = []
chapter_length = []
cnt = 0

chapter_titles.append(texts[0])
for i in range(1, len(texts)):
    if texts[i].strip():
        if re.match(r'제\s*\d+\s*장\s*[^\n]*\n?', texts[i]):
            chapter_titles.append(texts[i])
            chapters.append(chapter)
            chapter = []
            chapter_length.append(cnt)
            cnt = 0
        else:
            chapter.append(texts[i])
            cnt += 1

def predict_emotion(sentence, model, tokenizer):
    model.eval() # 평가 모드로 변경

    encoding = tokenizer(
        sentence,
        return_tensors='pt',
        padding='max_length',
        truncation=True,
        max_length=64
    )

    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)
    token_type_ids = encoding.get('token_type_ids', torch.zeros_like(input_ids)).to(device)

    with torch.no_grad():
        output = model(input_ids, attention_mask, token_type_ids)
        prediction = torch.argmax(output, dim=1).item()

    label_map_reverse = {
        0: '기쁨', 1: '슬픔', 2: '분노', 3: '불안', 4: '사랑', 5: '놀람'
    }

    return label_map_reverse[prediction]

import pandas as pd
from tqdm import tqdm

for i, chapter in enumerate(zip(chapter_titles, chapters)):
    title = chapter[0]
    sentences = chapter[1]
    results = []

    for sentence in tqdm(sentences, desc=f"'{title}' 감정 분석 중"):
        emotion = predict_emotion(sentence, model, tokenizer)
        results.append({
            "sentence": sentence,
            "emotion": emotion
        })

    # DataFrame으로 변환
    df = pd.DataFrame(results)

    # CSV 파일로 저장
    filename = f"{title}.csv"
    df.to_csv(filename, index=False, encoding='utf-8-sig')  # 윈도우 한글 호환
    print(f"저장 완료: {filename}")

# 사용할 감정 라벨
emotion_labels = ['기쁨', '슬픔', '분노', '불안', '놀람', '사랑']

# 각 감정의 극성 점수 (valence)
emotion_score = {
    '기쁨': 3, '사랑': 2, '놀람': 1,
    '불안': -1, '슬픔': -2, '분노': -3
}

def emotion_frequency_flow(emotion_df):
    # 감정 흐름 시각화
    plt.figure(figsize=(12, 6))

    for label in emotion_labels:
        x = [f'{i+1}장' for i in range(len(emotion_df))]
        y = emotion_df[label]
        plt.plot(x, y, marker='o', label=label)

    plt.title("챕터별 감정 빈도수")
    plt.xlabel("챕터")
    plt.ylabel("감정별 문장 갯수")
    plt.xticks(rotation=45)

    plt.legend()
    plt.tight_layout()
    plt.show()

# 감정 극성 평균값을 기반으로 흐름만 시각화
def emotion_avg_score_flow(emotion_df):
    plt.figure(figsize=(12, 6))

    x = [f'{i+1}장' for i in range(len(emotion_df))]
    y = emotion_df['avg_emotion_score']

    plt.plot(x, y, marker='o', color='darkblue', label='평균 감정 점수')

    plt.title("감정 극성 점수의 평균치 흐름")
    plt.xlabel("챕터")
    plt.ylabel("평균 감정 점수")
    plt.axhline(0, color='gray', linestyle='--', linewidth=1)
    plt.xticks(rotation=45)

    plt.legend()
    plt.tight_layout()
    plt.show()

# 평균 감정 점수의 분포를 히스토그램으로 시각화
def emotion_score_histogram(emotion_df):
    plt.figure(figsize=(12, 6))

    plt.hist(emotion_df['avg_emotion_score'], bins=10, color='steelblue', edgecolor='black')

    plt.title("평균 감정 점수 분포")
    plt.xlabel("평균 감정 점수")
    plt.ylabel("챕터 갯수")
    plt.axvline(0, color='gray', linestyle='--', linewidth=1)

    plt.tight_layout()
    plt.show()

# 감정별로 가장 우세한 챕터 상위 3개를 dictionary 형태로 반환

def get_chapter_dominance_by_emotion_dict(emotion_df, top_n=3):
    emotion_labels = ['기쁨', '슬픔', '분노', '불안', '놀람', '사랑']
    emotion_totals = {emotion: emotion_df[emotion].sum() for emotion in emotion_labels}

    dominance_dict = {}

    for emotion in emotion_labels:
        dominance_rows = []
        for _, row in emotion_df.iterrows():
            chapter = row['chapter']
            value = row[emotion]
            dominance_ratio = value / emotion_totals[emotion] if emotion_totals[emotion] > 0 else 0
            dominance_rows.append((chapter, round(dominance_ratio, 3)))

        top_chapters = sorted(dominance_rows, key=lambda x: x[1], reverse=True)[:top_n]
        dominance_dict[emotion] = top_chapters

    return dominance_dict

def analyze_emotion_structure(folder_path):
    chapter_emotions = []

    # 각 챕터 파일에서 감정 통계 집계
    for filename in sorted(os.listdir(folder_path)):
        if filename.endswith(".csv"):
            filepath = os.path.join(folder_path, filename)
            df = pd.read_csv(filepath)
            chapter_name = filename.replace(".csv", "")

            # 감정 빈도 수 계산
            emotion_counts = df['emotion'].value_counts().to_dict()
            row = {'chapter': chapter_name}
            total = 0
            score_sum = 0

            for label in emotion_labels:
                count = emotion_counts.get(label, 0)
                row[label] = count
                total += count
                score_sum += emotion_score.get(label, 0) * count

            row['total'] = total
            row['avg_emotion_score'] = score_sum / total if total > 0 else 0
            chapter_emotions.append(row)

    # 데이터프레임 생성
    emotion_df = pd.DataFrame(chapter_emotions)
    emotion_df = emotion_df.sort_values('chapter').reset_index(drop=True)
    emotion_df.fillna(0, inplace=True)

    # 그래프 시각화
    emotion_avg_score_flow(emotion_df)
    emotion_frequency_flow(emotion_df)
    emotion_score_histogram(emotion_df)

    # 특정 감정 많은 챕터 추출
    top_chapters_by_emotion = get_chapter_dominance_by_emotion_dict(emotion_df)

    return emotion_df, top_chapters_by_emotion

for (path, dirs, files) in os.walk(root_path):
    for dir in sorted(dirs):
        folder_path = path + '/' + dir  # 각 챕터 감정 CSV 폴더 경로
        emotion_df, top_chapters_by_emotion = analyze_emotion_structure(folder_path)
        # 감정 구조 분석 결과
        print(emotion_df.head())
        for label in emotion_labels:
            print(f"{label} 감정이 집중된 챕터 Top 3:")
            for k, v in top_chapters_by_emotion[label]:
                print(f"{k} 챕터: {v}")
            print()
    break