analyzer.py

"""
Utilitário de análise de genes para uso standalone
Pode ser importado como módulo ou executado como script
"""

import Bio
from Bio import SeqIO
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    classification_report,
)
import json


class ArabidopsisAnalyzer:
    """Classe para análise de genes do Arabidopsis thaliana"""

    def __init__(self, fasta_file):
        """Inicializa o analisador com um arquivo FASTA"""
        self.fasta_file = fasta_file
        self.sequences = []
        self.seq_ids = []
        self.features_df = None
        self.model = None
        self.scaler = None

    def load_sequences(self):
        """Carrega sequências do arquivo FASTA"""
        print("Carregando sequências...")
        for record in SeqIO.parse(self.fasta_file, "fasta"):
            self.sequences.append(str(record.seq).upper())
            self.seq_ids.append(record.id)
        print(f"✓ {len(self.sequences)} sequências carregadas")
        return len(self.sequences)

    def count_nucleotides(self, sequence):
        """Conta nucleotídeos em uma sequência"""
        return {
            "A": sequence.count("A"),
            "T": sequence.count("T"),
            "C": sequence.count("C"),
            "G": sequence.count("G"),
            "N": sequence.count("N"),
        }

    def calculate_gc_content(self, sequence):
        """Calcula conteúdo GC (%)"""
        gc_count = sequence.count("G") + sequence.count("C")
        return (gc_count / len(sequence) * 100) if len(sequence) > 0 else 0

    def extract_features(self):
        """Extrai features para machine learning"""
        print("Extraindo features...")
        features_list = []

        for seq in self.sequences:
            counts = self.count_nucleotides(seq)
            total = len(seq)

            features = {
                "length": total,
                "A_count": counts["A"],
                "T_count": counts["T"],
                "C_count": counts["C"],
                "G_count": counts["G"],
                "N_count": counts["N"],
                "A_freq": counts["A"] / total if total > 0 else 0,
                "T_freq": counts["T"] / total if total > 0 else 0,
                "C_freq": counts["C"] / total if total > 0 else 0,
                "G_freq": counts["G"] / total if total > 0 else 0,
                "GC_content": self.calculate_gc_content(seq),
                "AT_content": (
                    ((counts["A"] + counts["T"]) / total * 100) if total > 0 else 0
                ),
            }
            features_list.append(features)

        self.features_df = pd.DataFrame(features_list)
        print(f"✓ {len(self.features_df)} features extraídas")
        return self.features_df

    def get_summary_stats(self):
        """Retorna estatísticas resumidas"""
        if self.features_df is None:
            return None

        stats = {
            "Total de Sequências": len(self.sequences),
            "Comprimento Médio (bp)": self.features_df["length"].mean(),
            "Comprimento Máximo (bp)": self.features_df["length"].max(),
            "Comprimento Mínimo (bp)": self.features_df["length"].min(),
            "GC Content Médio (%)": self.features_df["GC_content"].mean(),
            "AT Content Médio (%)": self.features_df["AT_content"].mean(),
        }
        return stats

    def train_model(self, test_size=0.2, n_estimators=100, max_depth=10):
        """Treina um modelo de Random Forest"""
        print("\nTreinando modelo...")

        feature_cols = [
            "length",
            "A_freq",
            "T_freq",
            "C_freq",
            "G_freq",
            "GC_content",
            "AT_content",
        ]
        X = self.features_df[feature_cols] # type: ignore

        # Normalizar dados
        self.scaler = StandardScaler()
        X_scaled = self.scaler.fit_transform(X)

        # Classificação: GC content alto/baixo
        y = (
            self.features_df["GC_content"] > self.features_df["GC_content"].median() # type: ignore
        ).astype(int)

        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=test_size, random_state=42
        )

        self.model = RandomForestClassifier(
            n_estimators=n_estimators, max_depth=max_depth, random_state=42, n_jobs=-1
        )
        self.model.fit(X_train, y_train)

        # Avaliação
        y_pred = self.model.predict(X_test)

        metrics = {
            "Acurácia": accuracy_score(y_test, y_pred),
            "Precisão": precision_score(y_test, y_pred),
            "Recall": recall_score(y_test, y_pred),
            "F1-Score": f1_score(y_test, y_pred),
        }

        print("✓ Modelo treinado com sucesso!")
        for metric, value in metrics.items():
            print(f"  {metric}: {value:.4f}")

        return metrics

    def export_stats(self, filename="analysis_results.json"):
        """Exporga estatísticas em formato JSON"""
        stats = self.get_summary_stats()

        data = {
            "summary_statistics": stats,
            "features_summary": self.features_df.describe().to_dict(), # type: ignore
        }

        with open(filename, "w") as f:
            json.dump(data, f, indent=2)

        print(f"✓ Resultados exportados para {filename}")


def main():
    """Execução standalone"""
    fasta_file = "Arabidopsis_thaliana.TAIR10.dna.chromosome.1.fa"

    analyzer = ArabidopsisAnalyzer(fasta_file)
    analyzer.load_sequences()
    analyzer.extract_features()

    print("\n" + "=" * 50)
    print("ESTATÍSTICAS RESUMIDAS")
    print("=" * 50)
    stats = analyzer.get_summary_stats()
    for stat, value in stats.items(): # type: ignore
        if isinstance(value, float):
            print(f"{stat}: {value:.2f}")
        else:
            print(f"{stat}: {value}")

    print("\n" + "=" * 50)
    print("TREINAMENTO DO MODELO")
    print("=" * 50)
    metrics = analyzer.train_model()

    print("\n" + "=" * 50)
    print("EXPORTANDO RESULTADOS")
    print("=" * 50)
    analyzer.export_stats()


if __name__ == "__main__":
    main()