# models/regtech_nlp.py
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class ComplianceNLP:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("nlpaueb/legal-bert-base-uncased")
        self.model = AutoModelForSequenceClassification.from_pretrained(
            "nlpaueb/legal-bert-base-uncased", num_labels=2
        )

    def extract_obligations(self, text: str) -> list[dict]:
        sentences = self._split_sentences(text)
        obligations = []
        for sent in sentences:
            inputs = self.tokenizer(sent, return_tensors="pt", truncation=True, max_length=512)
            with torch.no_grad():
                outputs = self.model(**inputs)
            prob = torch.softmax(outputs.logits, dim=-1)[0][1].item()
            if prob > 0.7:
                obligations.append({
                    "text": sent,
                    "confidence": prob,
                    "type": self._classify_obligation(sent),
                })
        return obligations

    def _split_sentences(self, text: str) -> list[str]:
        import re
        return [s.strip() for s in re.split(r'[.!?]+', text) if len(s.strip()) > 20]

    def _classify_obligation(self, text: str) -> str:
        text_lower = text.lower()
        if any(w in text_lower for w in ["must", "shall", "required"]):
            return "mandatory"
        elif any(w in text_lower for w in ["should", "recommended"]):
            return "recommended"
        return "advisory"

# models/gap_analysis.py
import numpy as np
from sentence_transformers import SentenceTransformer

class GapAnalyzer:
    def __init__(self):
        self.model = SentenceTransformer("all-MiniLM-L6-v2")

    def analyze(self, regulations: list[str], policies: list[str]) -> dict:
        reg_embeddings = self.model.encode(regulations)
        policy_embeddings = self.model.encode(policies)

        gaps = []
        for i, reg in enumerate(regulations):
            similarities = [
                np.dot(reg_embeddings[i], policy_embeddings[j]) /
                (np.linalg.norm(reg_embeddings[i]) * np.linalg.norm(policy_embeddings[j]))
                for j in range(len(policies))
            ]
            max_sim = max(similarities) if similarities else 0
            if max_sim < 0.6:
                gaps.append({
                    "regulation": reg,
                    "max_similarity": max_sim,
                    "status": "uncovered",
                })
            elif max_sim < 0.8:
                gaps.append({
                    "regulation": reg,
                    "max_similarity": max_sim,
                    "status": "partial",
                })

        return {
            "total_regulations": len(regulations),
            "total_policies": len(policies),
            "covered": len([g for g in gaps if g["status"] == "covered"]),
            "partial": len([g for g in gaps if g["status"] == "partial"]),
            "uncovered": len([g for g in gaps if g["status"] == "uncovered"]),
            "gaps": gaps,
        }

# main.py
from models.regtech_nlp import ComplianceNLP
from models.gap_analysis import GapAnalyzer

nlp = ComplianceNLP()
analyzer = GapAnalyzer()

regulation_text = """
Financial institutions must implement customer due diligence procedures.
They shall verify customer identity using reliable, independent sources.
Suspicious activity must be reported to FinCEN within 30 days.
"""

obligations = nlp.extract_obligations(regulation_text)
print(f"Found {len(obligations)} obligations:")
for obs in obligations:
    print(f"  [{obs['type']}] {obs['text'][:80]}... (conf: {obs['confidence']:.2f})")

result = analyzer.analyze(
    regulations=["KYC verification required", "SAR filing within 30 days"],
    policies=["Customer identity verification process", "Transaction monitoring procedures"]
)
print(f"\nGap Analysis: {result['covered']} covered, {result['uncovered']} uncovered")