penracourses/research/views.py

import csv
import json
import random
from django.shortcuts import render, redirect, get_object_or_404
from django.contrib.auth.decorators import login_required
from django.core.exceptions import PermissionDenied
from django.db import transaction
from django.http import HttpResponse, JsonResponse, Http404
from django.urls import reverse
from django.utils import timezone
from django.utils.text import slugify
from django.utils.crypto import get_random_string
from django.contrib import messages

from atlas.models import CaseCollection, CaseDetail, CidReportAnswer
from generic.models import CidUser

from research.models import ResearchStudy, ResearchStudyArm, ResearchParticipant
from research.forms import (
    ResearchStudyForm,
    ResearchStudyArmForm,
    ResearchParticipantIntakeForm,
    BulkParticipantGenerationForm,
)


# ============================================================================
# Utility Functions
# ============================================================================

def _user_can_manage_research_study(study: ResearchStudy, user) -> bool:
    """Check if user can manage a research study."""
    if not user.is_authenticated:
        return False
    if user.is_superuser:
        return True
    if user.groups.filter(name="atlas_editor").exists():
        return True
    return study.author.filter(pk=user.pk).exists()


def _get_next_cid():
    """Get the next available CID number."""
    from django.db.models import Max
    max_cid = CidUser.objects.aggregate(Max('cid'))['cid__max'] or 0
    return max(max_cid + 1, 100000)


def _choose_research_arm(study: ResearchStudy, candidate_group: str = "") -> ResearchStudyArm:
    """Select an arm for participant assignment based on study's randomisation strategy."""
    arms = list(study.arms.filter(active=True).select_related("packet").order_by("sort_order", "id"))
    if not arms:
        raise Http404("No active packets configured for this study")

    randomisation_mode = study.randomisation_mode

    # Completely random: weighted by allocation_weight
    if randomisation_mode == ResearchStudy.RandomisationMode.COMPLETELY_RANDOM:
        weights = [max(1, a.allocation_weight) for a in arms]
        return random.choices(arms, weights=weights, k=1)[0]

    # Target-based: allocate until arm reaches target, then move to next
    if randomisation_mode == ResearchStudy.RandomisationMode.TARGET_BASED:
        targets = study.allocation_targets or {}
        eligible = []
        
        for arm in arms:
            target = targets.get(str(arm.pk), 0)
            if target <= 0:  # No limit
                eligible.append(arm)
            else:
                current_count = ResearchParticipant.objects.filter(
                    study=study,
                    assigned_arm=arm,
                ).count()
                if current_count < target:
                    eligible.append(arm)
        
        if not eligible:
            # All targets met; fall back to balanced
            participant_qs = ResearchParticipant.objects.filter(study=study, assigned_arm__isnull=False)
            if study.balance_within_candidate_group and candidate_group:
                participant_qs = participant_qs.filter(candidate_group=candidate_group)
            counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
            min_count = min(counts.values())
            eligible = [arm for arm in arms if counts[arm.pk] == min_count]
        
        return random.choice(eligible)

    # Balanced within group: balance per candidate_group
    if randomisation_mode == ResearchStudy.RandomisationMode.BALANCED_WITHIN_GROUP and candidate_group:
        participant_qs = ResearchParticipant.objects.filter(
            study=study,
            assigned_arm__isnull=False,
            candidate_group=candidate_group,
        )
        counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
        min_count = min(counts.values())
        eligible = [arm for arm in arms if counts[arm.pk] == min_count]
        return random.choice(eligible)

    # Balanced (default): global balance across all participants
    participant_qs = ResearchParticipant.objects.filter(study=study, assigned_arm__isnull=False)
    counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
    min_count = min(counts.values())
    eligible = [arm for arm in arms if counts[arm.pk] == min_count]
    return random.choice(eligible)


def _ensure_research_participant_cid(participant: ResearchParticipant) -> CidUser:
    """Ensure participant has a CidUser (creates if needed); maintains 1-1 mapping."""
    if participant.cid_user is not None:
        return participant.cid_user

    # Try to allocate a unique CID
    created_cid = None
    for _ in range(10):
        try:
            try:
                next_cid = _get_next_cid()
            except Exception:
                next_cid = random.randint(100000, 999999)
            
            created_cid = CidUser.objects.create(
                cid=next_cid,
                passcode=get_random_string(8),
                name=f"study:{participant.study.slug}:{participant.pseudo_id}",
                active=True,
            )
            break
        except Exception:
            continue

    if created_cid is None:
        raise RuntimeError("Unable to allocate CID user for participant")

    # Update participant with CidUser (1-1 relationship)
    try:
        participant.cid_user = created_cid
        participant.save(update_fields=["cid_user", "updated_at"])
    except Exception as e:
        # If 1-1 constraint violated, another participant may have been created; clean up
        created_cid.delete()
        raise RuntimeError(f"Unable to assign CID to participant: {e}")

    return created_cid


def _build_study_export_rows(study: ResearchStudy):
    """Build export rows for study results."""
    rows = []
    participants = (
        study.participants.select_related("assigned_arm", "assigned_arm__packet", "cid_user", "user_user")
        .order_by("pseudo_id")
    )

    for participant in participants:
        raw_score = None
        max_possible_score = None
        normalised_percent = None
        marked_answers = 0
        total_cases = 0
        attempt_start = None
        attempt_end = None
        attempt_completed = False

        if participant.assigned_arm and participant.cid_user:
            packet = participant.assigned_arm.packet
            casedetails = CaseDetail.objects.filter(collection=packet)
            total_cases = casedetails.count()
            answers = CidReportAnswer.objects.filter(
                question__in=casedetails,
                cid=participant.cid_user.cid,
            )
            marked_scores = [a.score for a in answers if a.score is not None]
            raw_score = sum(marked_scores) if marked_scores else 0
            marked_answers = len(marked_scores)

            # Get case max score from packet's marking config
            case_max = 10  # Default
            if packet.case_question_max_score:
                case_max = packet.case_question_max_score

            max_possible_score = total_cases * case_max
            if max_possible_score > 0:
                normalised_percent = round((raw_score / max_possible_score) * 100, 2)

            attempt = packet.get_cid_exams(cid_user=participant.cid_user).first()
            if attempt:
                attempt_start = attempt.start_time
                attempt_end = attempt.end_time
                attempt_completed = bool(attempt.completed)

        rows.append(
            {
                "study_id": study.pk,
                "study_slug": study.slug,
                "study_name": study.name,
                "pseudo_id": participant.pseudo_id,
                "candidate_group": participant.candidate_group,
                "age_band": participant.age_band,
                "sex": participant.sex,
                "training_grade": participant.training_grade,
                "years_experience": participant.years_experience,
                "email": participant.email,
                "consented": participant.consented,
                "assigned_arm": participant.assigned_arm.name if participant.assigned_arm else "",
                "packet": participant.assigned_arm.packet.name if participant.assigned_arm else "",
                "assignment_method": participant.assignment_method,
                "assigned_at": participant.assigned_at,
                "cid": participant.cid_user.cid if participant.cid_user else "",
                "passcode": participant.cid_user.passcode if participant.cid_user else "",
                "attempt_start": attempt_start,
                "attempt_end": attempt_end,
                "attempt_completed": attempt_completed,
                "total_cases": total_cases,
                "marked_answers": marked_answers,
                "raw_score": raw_score,
                "max_possible_score": max_possible_score,
                "normalised_percent": normalised_percent,
            }
        )

    return rows


# ============================================================================
# Study Management Views
# ============================================================================

@login_required
def study_list(request):
    """List research studies accessible to the user."""
    studies = ResearchStudy.objects.all().prefetch_related("author")
    if not (request.user.is_superuser or request.user.groups.filter(name="atlas_editor").exists()):
        studies = studies.filter(author=request.user)
    return render(request, "research/study_list.html", {"studies": studies.distinct()})


@login_required
def study_create(request):
    """Create a new research study."""
    if request.method == "POST":
        form = ResearchStudyForm(request.POST)
        if form.is_valid():
            study = form.save()
            study.author.add(request.user)
            messages.success(request, f"Study '{study.name}' created successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyForm()
    return render(request, "research/study_form.html", {"form": form, "is_create": True})


@login_required
def study_update(request, pk: int):
    """Update an existing research study."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    if request.method == "POST":
        form = ResearchStudyForm(request.POST, instance=study)
        if form.is_valid():
            form.save()
            messages.success(request, f"Study '{study.name}' updated successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyForm(instance=study)

    return render(
        request,
        "research/study_form.html",
        {"form": form, "study": study, "is_create": False},
    )


@login_required
def study_detail(request, pk: int):
    """Display study detail with arms and participants."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    arm_form = ResearchStudyArmForm() if request.method != "POST" else None
    if request.method == "POST" and "add_arm" in request.POST:
        arm_form = ResearchStudyArmForm(request.POST)
        if arm_form.is_valid():
            arm = arm_form.save(commit=False)
            arm.study = study
            arm.save()
            messages.success(request, f"Arm '{arm.name}' added successfully.")
            return redirect("research:study_detail", pk=study.pk)

    arms = study.arms.select_related("packet").all()
    participants = study.participants.select_related("assigned_arm", "assigned_arm__packet", "cid_user").all()
    arm_rows = [
        (arm, participants.filter(assigned_arm=arm).count())
        for arm in arms
    ]
    return render(
        request,
        "research/study_detail.html",
        {
            "study": study,
            "arms": arms,
            "participants": participants,
            "arm_rows": arm_rows,
            "arm_form": arm_form,
        },
    )


@login_required
def study_arm_update(request, pk: int, arm_id: int):
    """Update a study arm."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    arm = get_object_or_404(ResearchStudyArm, pk=arm_id, study=study)
    if request.method == "POST":
        form = ResearchStudyArmForm(request.POST, instance=arm, study=study)
        if form.is_valid():
            form.save()
            messages.success(request, f"Arm '{arm.name}' updated successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyArmForm(instance=arm, study=study)

    return render(
        request,
        "research/study_arm_form.html",
        {"study": study, "form": form, "arm": arm},
    )


@login_required
def study_export_csv(request, pk: int):
    """Export study results to CSV."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    rows = _build_study_export_rows(study)
    response = HttpResponse(content_type="text/csv")
    response["Content-Disposition"] = f'attachment; filename="study_{study.slug}_results.csv"'

    if rows:
        writer = csv.DictWriter(response, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(row)

    return response


@login_required
def study_export_json(request, pk: int):
    """Export study results to JSON."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    rows = _build_study_export_rows(study)
    return JsonResponse(rows, safe=False)


@login_required
def study_bulk_generate(request, pk: int):
    """Generate participants in bulk from CSV/JSON."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    generated_count = 0
    error_message = None

    if request.method == "POST":
        form = BulkParticipantGenerationForm(request.POST, request.FILES)
        if form.is_valid():
            try:
                upload_format = form.cleaned_data["format"]
                uploaded_file = request.FILES["file"]
                
                if upload_format == "csv":
                    # Parse CSV
                    decoded_file = uploaded_file.read().decode('utf-8').splitlines()
                    reader = csv.DictReader(decoded_file)
                    for row_num, row in enumerate(reader, start=2):  # +2 for header + 1-indexing
                        try:
                            pseudo_id = row.get("pseudo_id", "").strip()
                            candidate_group = row.get("candidate_group", "").strip()
                            email = row.get("email", "").strip()
                            
                            if not pseudo_id:
                                error_message = f"Row {row_num}: pseudo_id is required"
                                break
                            
                            participant, created = ResearchParticipant.objects.get_or_create(
                                study=study,
                                pseudo_id=pseudo_id,
                                defaults={
                                    "candidate_group": candidate_group,
                                    "email": email,
                                }
                            )
                            if created:
                                generated_count += 1
                        except Exception as e:
                            error_message = f"Row {row_num}: {str(e)}"
                            break

                elif upload_format == "json":
                    # Parse JSON
                    data = json.load(uploaded_file)
                    if not isinstance(data, list):
                        error_message = "JSON must be an array of participant objects"
                    else:
                        for row_num, row in enumerate(data, start=1):
                            try:
                                pseudo_id = row.get("pseudo_id", "").strip()
                                candidate_group = row.get("candidate_group", "").strip()
                                email = row.get("email", "").strip()
                                
                                if not pseudo_id:
                                    error_message = f"Record {row_num}: pseudo_id is required"
                                    break
                                
                                participant, created = ResearchParticipant.objects.get_or_create(
                                    study=study,
                                    pseudo_id=pseudo_id,
                                    defaults={
                                        "candidate_group": candidate_group,
                                        "email": email,
                                    }
                                )
                                if created:
                                    generated_count += 1
                            except Exception as e:
                                error_message = f"Record {row_num}: {str(e)}"
                                break

                if error_message:
                    messages.error(request, error_message)
                else:
                    messages.success(request, f"Generated {generated_count} new participants")
                    return redirect("research:study_detail", pk=study.pk)

            except Exception as e:
                error_message = f"File processing error: {str(e)}"
                messages.error(request, error_message)
    else:
        form = BulkParticipantGenerationForm()

    return render(
        request,
        "research/bulk_generate.html",
        {"study": study, "form": form},
    )


# ============================================================================
# Participant Entry & Access
# ============================================================================

def participant_entry(request, slug: str, pseudo_id: str):
    """Participant landing page with on-site randomisation and intake form."""
    study = get_object_or_404(ResearchStudy, slug=slug)
    if not study.active and not (request.user.is_authenticated and _user_can_manage_research_study(study, request.user)):
        raise Http404("Study not available")

    # Check generation mode
    if study.generation_mode == ResearchStudy.GenerationMode.BULK_ONLY:
        if not request.user.is_authenticated or not _user_can_manage_research_study(study, request.user):
            # Try to get existing participant
            try:
                participant = ResearchParticipant.objects.get(study=study, pseudo_id=pseudo_id)
            except ResearchParticipant.DoesNotExist:
                raise Http404("Participant not found. Contact study administrator.")
    else:
        # AUTO_ON_ACCESS: create if doesn't exist
        participant, _ = ResearchParticipant.objects.get_or_create(
            study=study,
            pseudo_id=pseudo_id,
        )

    if request.user.is_authenticated and participant.user_user is None:
        participant.user_user = request.user

    if request.method == "POST":
        form = ResearchParticipantIntakeForm(
            request.POST,
            instance=participant,
            collect_demographics=study.collect_demographics,
        )
        if form.is_valid():
            participant = form.save(commit=False)
            participant.study = study

            with transaction.atomic():
                if participant.assigned_arm is None:
                    assigned_arm = _choose_research_arm(
                        study,
                        candidate_group=participant.candidate_group,
                    )
                    participant.assigned_arm = assigned_arm
                    participant.assignment_method = study.randomisation_mode
                    participant.assigned_at = timezone.now()

                participant.save()
                cid_user = _ensure_research_participant_cid(participant)
                
                # Store email on CidUser if provided
                if participant.email and not cid_user.email:
                    cid_user.email = participant.email
                    cid_user.save(update_fields=['email'])

                # Ensure packet can be accessed via existing CID flow
                if participant.assigned_arm is not None:
                    participant.assigned_arm.packet.valid_cid_users.add(cid_user)

            messages.success(request, "Intake complete. You can now start the study.")
            return redirect(
                "research:participant_entry",
                slug=study.slug,
                pseudo_id=participant.pseudo_id,
            )
    else:
        form = ResearchParticipantIntakeForm(
            instance=participant,
            collect_demographics=study.collect_demographics,
        )

    # Build start URL if assigned
    assigned_packet = participant.assigned_arm.packet if participant.assigned_arm else None
    start_url = None
    if participant.assigned_arm and participant.cid_user:
        start_url = reverse(
            "atlas:collection_case_view_take",
            kwargs={
                "pk": assigned_packet.pk,
                "case_number": 0,
                "cid": participant.cid_user.cid,
                "passcode": participant.cid_user.passcode,
            },
        )

    return render(
        request,
        "research/participant_entry.html",
        {
            "study": study,
            "participant": participant,
            "form": form,
            "assigned_packet": assigned_packet,
            "start_url": start_url,
        },
    )