import streamlit as st
import geojson
from shapely.geometry import shape, mapping
from shapely.ops import unary_union
from transformers import pipeline
import geopandas as gpd
from difflib import get_close_matches
import re

# Initialize the LLM pipeline
classifier = pipeline(
    "text-classification",
    model="facebook/bart-large-mnli",
    return_all_scores=True
)

# Define command patterns and their variations
COMMAND_PATTERNS = {
    "fix invalid geometries": [
        r"fix.*invalid",
        r"repair.*invalid",
        r"correct.*invalid",
        r"invalid.*fix",
        r"geometry.*invalid",
        r"fix.*geometry",
        r"invalid.*repair",
        r"broken.*geometry",
        r"fix.*broken",
    ],
    "fix self intersections": [
        r"self.*intersect",
        r"intersect.*self",
        r"crossing.*lines",
        r"overlapping.*lines",
        r"fix.*intersect",
        r"repair.*intersect",
        r"intersections.*fix",
    ],
    "remove duplicate vertices": [
        r"duplicate.*vertex",
        r"duplicate.*vertices",
        r"remove.*duplicate",
        r"delete.*duplicate",
        r"double.*points",
        r"repeated.*vertex",
        r"duplicate.*points",
    ],
    "fix overlaps": [
        r"overlap.*fix",
        r"fix.*overlap",
        r"repair.*overlap",
        r"overlapping.*areas",
        r"overlapping.*polygons",
        r"merge.*overlap",
    ],
    "simplify geometry": [
        r"simplify",
        r"make.*simple",
        r"reduce.*vertices",
        r"smooth.*geometry",
        r"generalize",
        r"clean.*geometry",
    ],
    "validate geometry": [
        r"validate",
        r"check.*valid",
        r"verify.*geometry",
        r"geometry.*valid",
        r"check.*geometry",
    ]
}

def preprocess_text(text):
    # Common typo corrections
    typo_corrections = {
        'geomtry': 'geometry',
        'geometrie': 'geometry',
        'intersecton': 'intersection',
        'intersections': 'intersection',
        'vertice': 'vertex',
        'verticies': 'vertices',
        'duplicat': 'duplicate',
        'simplifie': 'simplify',
        'overlaping': 'overlapping',
        'repare': 'repair',
        'invalide': 'invalid'
    }

    words = text.lower().split()
    corrected_words = []

    for word in words:
        if word in typo_corrections:
            corrected_words.append(typo_corrections[word])
        else:
            # Find close matches in our vocabulary
            vocab = list(typo_corrections.keys()) + list(typo_corrections.values())
            matches = get_close_matches(word, vocab, n=1, cutoff=0.8)
            if matches:
                corrected_words.append(matches[0])
            else:
                corrected_words.append(word)

    return ' '.join(corrected_words)

def identify_command(text):
    text = preprocess_text(text)

    for command, patterns in COMMAND_PATTERNS.items():
        for pattern in patterns:
            if re.search(pattern, text, re.IGNORECASE):
                return command

    # If no direct match, use the LLM classifier
    results = classifier(text, list(COMMAND_PATTERNS.keys()))
    return max(results[0], key=lambda x: x['score'])['label']

def fix_geometry(geojson_data, command, tolerance=0.01):
    fixed_features = []
    stats = {
        'fixed_count': 0,
        'total_count': len(geojson_data['features']),
        'command': command
    }

    for feature in geojson_data['features']:
        original_geom = shape(feature['geometry'])
        fixed_geom = original_geom

        try:
            if command == "fix invalid geometries" or command == "fix self intersections":
                fixed_geom = original_geom.buffer(0)
                if not original_geom.equals(fixed_geom):
                    stats['fixed_count'] += 1

            elif command == "remove duplicate vertices":
                fixed_geom = original_geom.simplify(0, preserve_topology=True)
                if not original_geom.equals(fixed_geom):
                    stats['fixed_count'] += 1

            elif command == "fix overlaps":
                fixed_geom = unary_union([original_geom])
                if not original_geom.equals(fixed_geom):
                    stats['fixed_count'] += 1

            elif command == "simplify geometry":
                fixed_geom = original_geom.simplify(tolerance)
                if not original_geom.equals(fixed_geom):
                    stats['fixed_count'] += 1

            elif command == "validate geometry":
                if not original_geom.is_valid:
                    fixed_geom = original_geom.buffer(0)
                    stats['fixed_count'] += 1

            feature['geometry'] = mapping(fixed_geom)
            fixed_features.append(feature)

        except Exception as e:
            st.warning(f"Error processing feature: {str(e)}")
            fixed_features.append(feature)

    return {'type': 'FeatureCollection', 'features': fixed_features}, stats

def generate_response(stats):
    if stats['fixed_count'] == 0:
        return f"No geometries needed fixing for {stats['command']}"
    else:
        return f"Fixed {stats['fixed_count']} out of {stats['total_count']} geometries using {stats['command']}"

def main():
    st.title("Smart GeoJSON Geometry Fixer")
    st.write("Upload a GeoJSON file and use natural language commands to fix geometry issues.")

    # Add custom CSS for better chat display
    st.markdown("""
        <style>
        .chat-message {
            padding: 1rem;
            border-radius: 0.5rem;
            margin-bottom: 1rem;
        }
        .user-message {
            background-color: #e6f3ff;
        }
        .assistant-message {
            background-color: #f0f2f6;
        }
        </style>
    """, unsafe_allow_html=True)

    uploaded_file = st.file_uploader("Upload GeoJSON", type="geojson")

    if 'geojson_data' not in st.session_state and uploaded_file is not None:
        st.session_state.geojson_data = geojson.load(uploaded_file)

    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    if 'current_geojson' not in st.session_state and uploaded_file is not None:
        st.session_state.current_geojson = st.session_state.geojson_data

    # Command suggestions
    st.header("Example Commands")
    example_commands = [
        "Fix invalid geometries in this file",
        "Can you remove duplicate vertices?",
        "Please fix self intersections",
        "Check for overlapping areas",
        "Simplify the geometry",
        "Validate all geometries"
    ]
    for cmd in example_commands:
        st.write(f"- {cmd}")

    user_input = st.text_input("Enter your command:")

    if user_input:
        st.session_state.chat_history.append(("User", user_input))

        command = identify_command(user_input)

        if hasattr(st.session_state, 'current_geojson'):
            fixed_geojson, stats = fix_geometry(st.session_state.current_geojson, command)
            st.session_state.current_geojson = fixed_geojson

            response = generate_response(stats)
            st.session_state.chat_history.append(("Assistant", response))

            # Display the fixed GeoJSON
            st.subheader("Current GeoJSON State")
            st.json(fixed_geojson)

            # Download options
            col1, col2 = st.columns(2)
            with col1:
                st.download_button(
                    label="Download GeoJSON",
                    data=geojson.dumps(fixed_geojson),
                    file_name="fixed_geometry.geojson",
                    mime="application/geo+json"
                )
            with col2:
                if stats['fixed_count'] > 0:
                    st.success(f"✅ Fixed {stats['fixed_count']} issues")
                else:
                    st.info("✓ No issues found")

    # Display chat history with styling
    st.subheader("Chat History")
    for role, message in st.session_state.chat_history:
        css_class = "user-message" if role == "User" else "assistant-message"
        st.markdown(f'<div class="chat-message {css_class}"><b>{role}:</b> {message}</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()