"""
A Self Assessment module that allows students to write open-ended responses,
submit, then see a rubric and rate themselves.  Persists student supplied
hints, answers, and assessment judgment (currently only correct/incorrect).
Parses xml definition file--see below for exact format.
"""

import json
import logging
from lxml import etree
import capa.xqueue_interface as xqueue_interface

from xmodule.capa_module import ComplexEncoder
from xmodule.progress import Progress
from xmodule.stringify import stringify_children
from capa.util import *
import openendedchild

from numpy import median

from datetime import datetime
from pytz import UTC

from .combined_open_ended_rubric import CombinedOpenEndedRubric

log = logging.getLogger("edx.courseware")

class OpenEndedModule(openendedchild.OpenEndedChild):
    """
    The open ended module supports all external open ended grader problems.
    Sample XML file:
    <openended min_score_to_attempt="1" max_score_to_attempt="1">
        <openendedparam>
            <initial_display>Enter essay here.</initial_display>
            <answer_display>This is the answer.</answer_display>
            <grader_payload>{"grader_settings" : "ml_grading.conf", "problem_id" : "6.002x/Welcome/OETest"}</grader_payload>
        </openendedparam>
    </openended>
    """

    TEMPLATE_DIR = "combinedopenended/openended"

    def setup_response(self, system, location, definition, descriptor):
        """
        Sets up the response type.
        @param system: Modulesystem object
        @param location: The location of the problem
        @param definition: The xml definition of the problem
        @param descriptor: The OpenEndedDescriptor associated with this
        @return: None
        """
        oeparam = definition['oeparam']

        self.url = definition.get('url', None)
        self.queue_name = definition.get('queuename', self.DEFAULT_QUEUE)
        self.message_queue_name = definition.get('message-queuename', self.DEFAULT_MESSAGE_QUEUE)

        # This is needed to attach feedback to specific responses later
        self.submission_id = None
        self.grader_id = None

        error_message = "No {0} found in problem xml for open ended problem. Contact the learning sciences group for assistance."
        if oeparam is None:
            # This is a staff_facing_error
            raise ValueError(error_message.format('oeparam'))
        if self.child_prompt is None:
            raise ValueError(error_message.format('prompt'))
        if self.child_rubric is None:
            raise ValueError(error_message.format('rubric'))

        self._parse(oeparam, self.child_prompt, self.child_rubric, system)

        if self.child_created is True and self.child_state == self.ASSESSING:
            self.child_created = False
            self.send_to_grader(self.latest_answer(), system)
            self.child_created = False

    def _parse(self, oeparam, prompt, rubric, system):
        '''
        Parse OpenEndedResponse XML:
            self.initial_display
            self.payload - dict containing keys --
            'grader' : path to grader settings file, 'problem_id' : id of the problem

            self.answer - What to display when show answer is clicked
        '''
        # Note that OpenEndedResponse is agnostic to the specific contents of grader_payload
        prompt_string = stringify_children(prompt)
        rubric_string = stringify_children(rubric)
        self.child_prompt = prompt_string
        self.child_rubric = rubric_string

        grader_payload = oeparam.find('grader_payload')
        grader_payload = grader_payload.text if grader_payload is not None else ''

        # Update grader payload with student id.  If grader payload not json, error.
        try:
            parsed_grader_payload = json.loads(grader_payload)
            # NOTE: self.system.location is valid because the capa_module
            # __init__ adds it (easiest way to get problem location into
            # response types)
        except TypeError, ValueError:
            # This is a dev_facing_error
            log.exception(
                "Grader payload from external open ended grading server is not a json object! Object: {0}".format(
                    grader_payload))

        self.initial_display = find_with_default(oeparam, 'initial_display', '')
        self.answer = find_with_default(oeparam, 'answer_display', 'No answer given.')

        parsed_grader_payload.update({
            'location': self.location_string,
            'course_id': system.course_id,
            'prompt': prompt_string,
            'rubric': rubric_string,
            'initial_display': self.initial_display,
            'answer': self.answer,
            'problem_id': self.display_name,
            'skip_basic_checks': self.skip_basic_checks,
            'control': json.dumps(self.control),
        })
        updated_grader_payload = json.dumps(parsed_grader_payload)

        self.payload = {'grader_payload': updated_grader_payload}

    def skip_post_assessment(self, _data, system):
        """
        Ajax function that allows one to skip the post assessment phase
        @param data: AJAX dictionary
        @param system: ModuleSystem
        @return: Success indicator
        """
        self.child_state = self.DONE
        return {'success': True}

    def message_post(self, data, system):
        """
        Handles a student message post (a reaction to the grade they received from an open ended grader type)
        Returns a boolean success/fail and an error message
        """

        event_info = dict()
        event_info['problem_id'] = self.location_string
        event_info['student_id'] = system.anonymous_student_id
        event_info['survey_responses'] = data

        survey_responses = event_info['survey_responses']
        for tag in ['feedback', 'submission_id', 'grader_id', 'score']:
            if tag not in survey_responses:
                # This is a student_facing_error
                return {'success': False,
                        'msg': "Could not find needed tag {0} in the survey responses.  Please try submitting again.".format(
                            tag)}
        try:
            submission_id = int(survey_responses['submission_id'])
            grader_id = int(survey_responses['grader_id'])
            feedback = str(survey_responses['feedback'].encode('ascii', 'ignore'))
            score = int(survey_responses['score'])
        except:
            # This is a dev_facing_error
            error_message = (
                "Could not parse submission id, grader id, "
                "or feedback from message_post ajax call.  "
                "Here is the message data: {0}".format(survey_responses)
            )
            log.exception(error_message)
            # This is a student_facing_error
            return {'success': False, 'msg': "There was an error saving your feedback.  Please contact course staff."}

        xqueue = system.get('xqueue')
        if xqueue is None:
            return {'success': False, 'msg': "Couldn't submit feedback."}
        qinterface = xqueue['interface']
        qtime = datetime.strftime(datetime.now(UTC), xqueue_interface.dateformat)
        anonymous_student_id = system.anonymous_student_id
        queuekey = xqueue_interface.make_hashkey(str(system.seed) + qtime +
                                                 anonymous_student_id +
                                                 str(len(self.child_history)))

        xheader = xqueue_interface.make_xheader(
            lms_callback_url=xqueue['construct_callback'](),
            lms_key=queuekey,
            queue_name=self.message_queue_name
        )

        student_info = {
            'anonymous_student_id': anonymous_student_id,
            'submission_time': qtime,
        }
        contents = {
            'feedback': feedback,
            'submission_id': submission_id,
            'grader_id': grader_id,
            'score': score,
            'student_info': json.dumps(student_info),
        }

        (error, msg) = qinterface.send_to_queue(
            header=xheader,
            body=json.dumps(contents)
        )

        # Convert error to a success value
        success = True
        if error:
            success = False

        self.child_state = self.DONE

        # This is a student_facing_message
        return {'success': success, 'msg': "Successfully submitted your feedback."}

    def send_to_grader(self, submission, system):
        """
        Send a given submission to the grader, via the xqueue
        @param submission: The student submission to send to the grader
        @param system: Modulesystem
        @return: Boolean true (not useful right now)
        """

        # Prepare xqueue request
        #------------------------------------------------------------

        xqueue = system.get('xqueue')
        if xqueue is None:
            return False
        qinterface = xqueue['interface']
        qtime = datetime.strftime(datetime.now(UTC), xqueue_interface.dateformat)

        anonymous_student_id = system.anonymous_student_id

        # Generate header
        queuekey = xqueue_interface.make_hashkey(str(system.seed) + qtime +
                                                 anonymous_student_id +
                                                 str(len(self.child_history)))

        xheader = xqueue_interface.make_xheader(
            lms_callback_url=xqueue['construct_callback'](),
            lms_key=queuekey,
            queue_name=self.queue_name
        )

        contents = self.payload.copy()

        # Metadata related to the student submission revealed to the external grader
        student_info = {
            'anonymous_student_id': anonymous_student_id,
            'submission_time': qtime,
        }

        # Update contents with student response and student info
        contents.update({
            'student_info': json.dumps(student_info),
            'student_response': submission,
            'max_score': self.max_score(),
        })

        # Submit request. When successful, 'msg' is the prior length of the queue
        qinterface.send_to_queue(
            header=xheader,
            body=json.dumps(contents)
        )

        # State associated with the queueing request
        queuestate = {
            'key': queuekey,
            'time': qtime,
        }
        return True

    def _update_score(self, score_msg, queuekey, system):
        """
        Called by xqueue to update the score
        @param score_msg: The message from xqueue
        @param queuekey: The key sent by xqueue
        @param system: Modulesystem
        @return: Boolean True (not useful currently)
        """
        new_score_msg = self._parse_score_msg(score_msg, system)
        if not new_score_msg['valid']:
            new_score_msg['feedback'] = 'Invalid grader reply. Please contact the course staff.'

        # self.child_history is initialized as [].  record_latest_score() and record_latest_post_assessment()
        # operate on self.child_history[-1].  Thus we have to make sure child_history is not [].
        # Handle at this level instead of in record_*() because this is a good place to reduce the number of conditions
        # and also keep the persistent state from changing.
        if self.child_history:
            self.record_latest_score(new_score_msg['score'])
            self.record_latest_post_assessment(score_msg)
            self.child_state = self.POST_ASSESSMENT
        else:
            log.error((
                "Trying to update score without existing studentmodule child_history:\n"
                "   location: {location}\n"
                "   score: {score}\n"
                "   grader_ids: {grader_ids}\n"
                "   submission_ids: {submission_ids}").format(
                    location=self.location_string,
                    score=new_score_msg['score'],
                    grader_ids=new_score_msg['grader_ids'],
                    submission_ids=new_score_msg['submission_ids']
                )
            )

        return True

    def get_answers(self):
        """
        Gets and shows the answer for this problem.
        @return: Answer html
        """
        anshtml = '
{0}
'.format(self.answer)
        return {self.answer_id: anshtml}

    def get_initial_display(self):
        """
        Gets and shows the initial display for the input box.
        @return: Initial display html
        """
        return {self.answer_id: self.initial_display}

    def _convert_longform_feedback_to_html(self, response_items):
        """
        Take in a dictionary, and return html strings for display to student.
        Input:
            response_items: Dictionary with keys success, feedback.
                if success is True, feedback should be a dictionary, with keys for
                   types of feedback, and the corresponding feedback values.
                if success is False, feedback is actually an error string.

                NOTE: this will need to change when we integrate peer grading, because
                that will have more complex feedback.

        Output:
            String -- html that can be displayincorrect-icon.pnged to the student.
        """

        # We want to display available feedback in a particular order.
        # This dictionary specifies which goes first--lower first.
        priorities = {
            # These go at the start of the feedback
            'spelling': 0,
            'grammar': 1,
            # needs to be after all the other feedback
            'markup_text': 3
        }
        do_not_render = ['topicality', 'prompt-overlap']

        default_priority = 2

        def get_priority(elt):
            """
            Args:
                elt: a tuple of feedback-type, feedback
            Returns:
                the priority for this feedback type
            """
            return priorities.get(elt[0], default_priority)

        def encode_values(feedback_type, value):
            feedback_type = str(feedback_type).encode('ascii', 'ignore')
            if not isinstance(value, basestring):
                value = str(value)
            value = value.encode('ascii', 'ignore')
            return feedback_type, value

        def format_feedback(feedback_type, value):
            feedback_type, value = encode_values(feedback_type, value)
            feedback = """

            {value}

            """.format(feedback_type=feedback_type, value=value)
            return feedback

        def format_feedback_hidden(feedback_type, value):
            feedback_type, value = encode_values(feedback_type, value)
            feedback = """
            <input class="{feedback_type}" type="hidden" value="{value}" />
            """.format(feedback_type=feedback_type, value=value)
            return feedback

        # TODO (vshnayder): design and document the details of this format so
        # that we can do proper escaping here (e.g. are the graders allowed to
        # include HTML?)

        for tag in ['success', 'feedback', 'submission_id', 'grader_id']:
            if tag not in response_items:
                # This is a student_facing_error
                return format_feedback('errors', 'Error getting feedback from grader.')

        feedback_items = response_items['feedback']
        try:
            feedback = json.loads(feedback_items)
        except (TypeError, ValueError):
            # This is a dev_facing_error
            log.exception("feedback_items from external open ended grader have invalid json {0}".format(feedback_items))
            # This is a student_facing_error
            return format_feedback('errors', 'Error getting feedback from grader.')

        if response_items['success']:
            if len(feedback) == 0:
                # This is a student_facing_error
                return format_feedback('errors', 'No feedback available from grader.')

            for tag in do_not_render:
                if tag in feedback:
                    feedback.pop(tag)

            feedback_lst = sorted(feedback.items(), key=get_priority)
            feedback_list_part1 = u"\n".join(format_feedback(k, v) for k, v in feedback_lst)
        else:
            # This is a student_facing_error
            feedback_list_part1 = format_feedback('errors', response_items['feedback'])

        feedback_list_part2 = (u"\n".join([format_feedback_hidden(feedback_type, value)
                                           for feedback_type, value in response_items.items()
                                           if feedback_type in ['submission_id', 'grader_id']]))

        return u"\n".join([feedback_list_part1, feedback_list_part2])

    def _format_feedback(self, response_items, system):
        """
        Input:
            Dictionary called feedback.  Must contain keys seen below.
        Output:
            Return error message or feedback template
        """

        rubric_feedback = ""
        feedback = self._convert_longform_feedback_to_html(response_items)
        rubric_scores = []
        if response_items['rubric_scores_complete'] is True:
            rubric_renderer = CombinedOpenEndedRubric(system, True)
            rubric_dict = rubric_renderer.render_rubric(response_items['rubric_xml'])
            success = rubric_dict['success']
            rubric_feedback = rubric_dict['html']
            rubric_scores = rubric_dict['rubric_scores']

        if not response_items['success']:
            return system.render_template(
                "{0}/open_ended_error.html".format(self.TEMPLATE_DIR),
                {'errors': feedback}
            )

        feedback_template = system.render_template("{0}/open_ended_feedback.html".format(self.TEMPLATE_DIR), {
            'grader_type': response_items['grader_type'],
            'score': "{0} / {1}".format(response_items['score'], self.max_score()),
            'feedback': feedback,
            'rubric_feedback': rubric_feedback
        })

        return feedback_template, rubric_scores

    def _parse_score_msg(self, score_msg, system, join_feedback=True):
        """
         Grader reply is a JSON-dump of the following dict
           { 'correct': True/False,
             'score': Numeric value (floating point is okay) to assign to answer
             'msg': grader_msg
             'feedback' : feedback from grader
             'grader_type': what type of grader resulted in this score
             'grader_id': id of the grader
             'submission_id' : id of the submission
             'success': whether or not this submission was successful
             'rubric_scores': a list of rubric scores
             'rubric_scores_complete': boolean if rubric scores are complete
             'rubric_xml': the xml of the rubric in string format
             }

        Returns (valid_score_msg, correct, score, msg):
            valid_score_msg: Flag indicating valid score_msg format (Boolean)
            correct:         Correctness of submission (Boolean)
            score:           Points to be assigned (numeric, can be float)
        """
        fail = {
            'valid': False,
            'score': 0,
            'feedback': '',
            'rubric_scores': [[0]],
            'grader_types': [''],
            'feedback_items': [''],
            'feedback_dicts': [{}],
            'grader_ids': [0],
            'submission_ids': [0],
        }
        try:
            score_result = json.loads(score_msg)
        except (TypeError, ValueError):
            # This is a dev_facing_error
            error_message = ("External open ended grader message should be a JSON-serialized dict."
                             " Received score_msg = {0}".format(score_msg))
            log.error(error_message)
            fail['feedback'] = error_message
            return fail

        if not isinstance(score_result, dict):
            # This is a dev_facing_error
            error_message = ("External open ended grader message should be a JSON-serialized dict."
                             " Received score_result = {0}".format(score_result))
            log.error(error_message)
            fail['feedback'] = error_message
            return fail

        for tag in ['score', 'feedback', 'grader_type', 'success', 'grader_id', 'submission_id']:
            if tag not in score_result:
                # This is a dev_facing_error
                error_message = ("External open ended grader message is missing required tag: {0}"
                                 .format(tag))
                log.error(error_message)
                fail['feedback'] = error_message
                return fail
                # This is to support peer grading
        if isinstance(score_result['score'], list):
            feedback_items = []
            rubric_scores = []
            grader_types = []
            feedback_dicts = []
            grader_ids = []
            submission_ids = []
            for i in xrange(0, len(score_result['score'])):
                new_score_result = {
                    'score': score_result['score'][i],
                    'feedback': score_result['feedback'][i],
                    'grader_type': score_result['grader_type'],
                    'success': score_result['success'],
                    'grader_id': score_result['grader_id'][i],
                    'submission_id': score_result['submission_id'],
                    'rubric_scores_complete': score_result['rubric_scores_complete'][i],
                    'rubric_xml': score_result['rubric_xml'][i],
                }
                feedback_template, rubric_score = self._format_feedback(new_score_result, system)
                feedback_items.append(feedback_template)
                rubric_scores.append(rubric_score)
                grader_types.append(score_result['grader_type'])
                try:
                    feedback_dict = json.loads(score_result['feedback'][i])
                except Exception:
                    feedback_dict = score_result['feedback'][i]
                feedback_dicts.append(feedback_dict)
                grader_ids.append(score_result['grader_id'][i])
                submission_ids.append(score_result['submission_id'])
            if join_feedback:
                feedback = "".join(feedback_items)
            else:
                feedback = feedback_items
            score = int(round(median(score_result['score'])))
        else:
            # This is for instructor and ML grading
            feedback, rubric_score = self._format_feedback(score_result, system)
            score = score_result['score']
            rubric_scores = [rubric_score]
            grader_types = [score_result['grader_type']]
            feedback_items = [feedback]
            try:
                feedback_dict = json.loads(score_result['feedback'])
            except Exception:
                feedback_dict = score_result.get('feedback', '')
            feedback_dicts = [feedback_dict]
            grader_ids = [score_result['grader_id']]
            submission_ids = [score_result['submission_id']]

        self.submission_id = score_result['submission_id']


#!/usr/bin/env python

#
#@BEGIN LICENSE
#
# PSI4: an ab initio quantum chemistry software package
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#@END LICENSE
#
from __future__ import absolute_import
from __future__ import print_function
import sys
import math
import os
import re
import importlib
import collections

qcdbpkg_path = os.path.dirname(__file__)
sys.path.append(qcdbpkg_path + '/../')
import qcdb
import qcdb.basislist
from qcdb.exceptions import *
sys.path.append(qcdbpkg_path + '/../databases')

# load docstring info from database files (doesn't actually import database modules)
DBdocstrings = qcdb.dictify_database_docstrings()

# instructions
print("""
 Welcome to imake-db.
    Just fill in the variables when prompted.
    Hit ENTER to accept default.
    Strings should not be in quotes.
    Elements in arrays should be space-delimited.
    Nothing is case sensitive.
""")

# query database name
module_choices = dict(zip([x.upper() for x in DBdocstrings.keys()], DBdocstrings.keys()))

print('\n Choose your database.')
for item in module_choices.keys():
    print("""    %-12s   %s""" % ('[' + module_choices[item] + ']', DBdocstrings[module_choices[item]]['general'][0].lstrip(' |')))
print('\n')

user_obedient = False
while not user_obedient:
    temp = raw_input('    dbse = ').strip()
    if temp.upper() in module_choices.keys():
        db_name = module_choices[temp.upper()]
        user_obedient = True

# query database subset
subset_choices = dict(zip([x.upper() for x in DBdocstrings[db_name]['subset'].keys()], DBdocstrings[db_name]['subset'].keys()))

print('\n Choose your subset (multiple allowed).')
for key, val in DBdocstrings[db_name]['subset'].items():
    print("""    %-12s   %s""" % ('[' + key + ']', val))
print('\n')

subset = []
user_obedient = False
while not user_obedient:
    temp = raw_input('    subset [all] = ').strip()
    ltemp = temp.split()
    if temp == "":
        user_obedient = True
    for item in ltemp:
        if item.upper() in subset_choices.keys():
            subset.append(subset_choices[item.upper()])
            user_obedient = True
        else:
            user_obedient = False
            subset = []
            break

# query qc program
print("""
 Choose your quantum chemistry program.
    [qchem]
    [molpro]       writes Molpro input files
    [molpro2]      writes Molpro input files
    [psi4]         writes Psi4 input files
    #[nwchem]
    #[xyz]          writes basic xyz files only
""")

user_obedient = False
while not user_obedient:
    temp = raw_input('    qcprog = ').strip()
    if temp.lower() in ['molpro', 'psi4', 'molpro2', 'qchem']:
        qcprog = temp.lower()
        user_obedient = True

# Load module for QC program
try:
    qcmod = importlib.import_module('qcdb.' + qcprog)
except ImportError:
    print('\nPython module for QC program %s failed to load\n\n' % (qcprog))
    print('\nSearch path that was tried:\n')
    print(", ".join(map(str, sys.path)))
    raise ValidationError("Python module loading problem for QC program " + str(qcprog))
else:
    print(qcmod)
    qcmtdIN = qcmod.qcmtdIN

# query quantum chemical method(s)
method_choices = dict(zip([x.upper() for x in qcmtdIN.keys()], qcmtdIN.keys()))

print('\n Choose your quantum chemical methods (multiple allowed).')
for key, val in qcmtdIN.items():
    print("""    %-12s""" % ('[' + key + ']'))
print('\n')

methods = []
user_obedient = False
while not user_obedient:
    temp = raw_input('    methods = ').strip()
    ltemp = temp.split()
    for item in ltemp:
        if item.upper() in method_choices:
            methods.append(method_choices[item.upper()])
            user_obedient = True
        else:
            user_obedient = False
            methods = []
            break

# query basis set(s)
print("""
 Choose your basis set (multiple allowed).
    e.g., aug-cc-pvdz or 6-31+G* or cc-pvtz may-cc-pvtz aug-cc-pvtz
""")

bases = []
user_obedient = False
while not user_obedient:
    temp = raw_input('    bases = ').strip()
    ltemp = temp.split()
    for item in ltemp:
        btemp = qcdb.basislist.corresponding_basis(item, role='BASIS')
        if btemp:
            bases.append(btemp)
            user_obedient = True
        else:
            print('    Basis set %s not recognized.' % (item))
            proceed = qcdb.query_yes_no('    Proceed anyway? =', False)
            if proceed:
                bases.append(item)
                user_obedient = True
            else:
                bases = []
                user_obedient = False
                break

# query castup preference
print("""
 Do cast up from smaller basis set?
""")

user_obedient = False
while not user_obedient:
    castup = qcdb.query_yes_no('    castup [F] = ', False)
    user_obedient = True
    # below, options['SCF']['BASIS_GUESS']['value'] = castup

# query directory prefix
print("""
 State your destination directory prefix.
""")

user_obedient = False
while not user_obedient:
    temp = raw_input('    dirprefix [try] = ').strip()
    if temp == "":
        dirprefix = 'try'
        user_obedient = True
    if temp.isalnum():
        dirprefix = temp
        user_obedient = True

# query memory
print("""
 Choose your memory usage in MB.
""")

user_obedient = False
while not user_obedient:
    temp = raw_input('    memory [1600] = ').strip()
    if temp == "":
        memory = 1600
        user_obedient = True
    if temp.isdigit():
        memory = int(temp)
        user_obedient = True

# Load module for requested database
try:
    database = __import__(db_name)
except ImportError:
    print('\nPython module for database %s failed to load\n\n' % (db_name))
    print('\nSearch path that was tried:\n')
    print(", ".join(map(str, sys.path)))
    raise ValidationError("Python module loading problem for database " + str(db_name))
else:
    dbse = database.dbse
    HRXN = database.HRXN
    ACTV = database.ACTV
    RXNM = database.RXNM
    BIND = database.BIND
    TAGL = database.TAGL
    GEOS = database.GEOS
    try:
        DATA = database.DATA
    except AttributeError:
        DATA = {}

print("""
        <<< SCANNED SETTINGS  SCANNED SETTINGS  SCANNED SETTINGS  SCANNED SETTINGS >>>

                          dbse = %s
                        subset = %s
                        qcprog = %s
                       methods = %s
                         bases = %s
                     dirprefix = %s


''' Blurs all weights in the selected vertex group.

This is still one of my first blender addons, so please keep in mind I may be
implementing various operations badly or just plain wrong.'''

'''
*******************************************************************************
	License and Copyright
	Copyright 2012 Jordan Hueckstaedt
	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''

bl_info = {
	'name': 'Blur Weights',
	'author': 'Jordan Hueckstaedt',
	'version': (1, 0),
	'blender': (2, 63, 0),
	'location': 'View > Weight Tools > Blur',
	'warning': '', # used for warning icon and text in addons panel
	'description': 'Blurs the weights in the selected vertex group.',
	"wiki_url": "",
	"tracker_url": "",
	"support": 'TESTING',
	'category': 'Paint'
}

import bpy
import bmesh
import math
# from mathutils import *

class BlurSettingsCollection(bpy.types.PropertyGroup):
	iterations = bpy.props.IntProperty(
		name = "Iter",
		description = "Iterate this many times",
		default = 1,
		min = 1,
		soft_max = 20)

	factor = bpy.props.FloatProperty(
		name = "Factor",
		description = "Laplace Factor",
		default = 0.5,
		soft_min = 0.0,
		soft_max = 1.0)

	operation = bpy.props.EnumProperty(
		name="Blur Operation",
		items =	(('GAUSSIAN', "Gaussian", "Gaussian Blur"),
				('AVERAGE', "Average", "Simple Average Blur"),
				),
		)

	blur_type = bpy.props.EnumProperty(
		name="Blur Type",
		items =	(('0', "Normal", "Normal Blur"),
				('1', "Shrink", "Blur will only decrease values"),
				('2', "Grow", "Blur will only increase values"),
				),
		)

	selected_only = bpy.props.BoolProperty(
		name = "Selected Only",
		default = False)

class BlurWeights( object ):
	def __init__(self, active_index = None):

		obj = bpy.context.active_object
		if active_index is None:
			active_index = obj.vertex_groups.active_index

		# Get a BMesh representation to access connectivity information
		bm_obj = bmesh.new()
		bm_obj.from_mesh(obj.data)

		# Find face masking.  Faces can only be found in bmesh, however,
		# in paint mode vert selection is not updated correctly to edit mode
		# or in bmesh, so vert selection can only be found through object.data
		# I suspect this is a bug.
		face_mask = obj.data.use_paint_mask
		vertex_mask = obj.data.use_paint_mask_vertex
		if face_mask:
			masked_face_verts = set([v.index for f in bm_obj.faces if f.select for v in f.verts])

		weights = [1.0 for x in range(len(obj.data.vertices))]
		vert_indexes = []
		vert_indexes_test = []
		vert_group_indexes = []
		connected_verts = []
		gaussian_weights = []

		# Convenience function to find group index
		def vert_group_index(vert, group_id):
			for x, group_info in enumerate(vert.groups):
				if group_info.group == group_id:
					return x
			return None

		# Get weight info.  It's weird, I know
		for vert in obj.data.vertices:
			i = vert_group_index(vert, active_index)
			if i is not None:
				weights[vert.index] = vert.groups[i].weight

				# Skip if vert is not in mask.
				if not vert.hide and (not vertex_mask or vert.select) and (not face_mask or vert.index in masked_face_verts):
					# Get group info
					vert_indexes.append(vert.index)
					vert_group_indexes.append(i)
				else:
					vert_indexes_test.append(vert.index)

		# Sets are MUCH faster to test in than lists.  This made __init__ run 3x faster for me.
		vert_indexes_test = set(vert_indexes + vert_indexes_test)

		# Only use connected verts in which are in the group.  God this part was a bitch to get working.
		# Also, pre-calculate the guassian weights
		for x in reversed(range(len(vert_indexes))):
			vert = bm_obj.verts[vert_indexes[x]]

			connected_vert = [v.index for edge in vert.link_edges for v in edge.verts if v.index != vert.index and v.index in vert_indexes_test]

			inclusive_verts = connected_vert + [vert.index]
			link_edges = [edge for edge in vert.link_edges if edge.verts[0].index in inclusive_verts and edge.verts[1].index in inclusive_verts]
			if not link_edges:
				# Orphaned vert
				vert_indexes.pop(x)
				vert_group_indexes.pop(x)
			else:
				connected_verts.append(connected_vert)

				# Pre-Calculate gaussian weights.
				avg_edge = sum([edge.calc_length() for edge in link_edges]) / len(link_edges)

				gaussian_weight = {'total_weight' : 0.0}
				for i in connected_vert:
					distance = (vert.co - bm_obj.verts[i].co).length_squared
					gaussian_weight[str(i)] = (1.0 /(avg_edge * math.sqrt(2.0 * math.pi))) * math.exp(-(distance)/(2.0 * avg_edge ** 2))
					gaussian_weight['total_weight'] += gaussian_weight[str(i)]

				gaussian_weights.append( gaussian_weight )

		# The last for loop was iterating in reverse to allow removing indexes while iterating.
		# So these lists are reversed
		gaussian_weights.reverse()
		connected_verts.reverse()

		self.weights = weights
		self.vert_indexes = vert_indexes
		self.vert_group_indexes = vert_group_indexes
		self.connected_verts = connected_verts
		self.gaussian_weights = gaussian_weights
		self.active_index = active_index
		bm_obj.free()

	def execute(self, iterations = 1, factor = 0.5, do_gaussian = True, blur_type = 0, selected_only = False):
		''' Calculate the blurred weights.
		blur_type options:
		0 - Normal
		1 - Shrink
		2 - Grow
		'''

		# Whatever you do, never save obj to a class member variable.  Blender will freeze.
		obj = bpy.context.active_object

		# Make a copy of weights.  Otherwise edits will be compounded onto the initial weight set.
		weights = self.weights[:]

		for i in range(iterations):
			new_weights = weights
			for x, i in enumerate(self.vert_indexes):
				if not selected_only  or (not do_gaussian and obj.data.vertices[i].select):
					# Skip guassian if the denominator in the weight function is 0
					# Which would most likely mean all the connected verts are in the same position
					if do_gaussian and self.gaussian_weights[x]['total_weight'] > 0.0:
						average_weight = 0
						for v in self.connected_verts[x]:
							average_weight += self.gaussian_weights[x][str(v)] / self.gaussian_weights[x]['total_weight'] * weights[v]
					else:
						# average_weight = sum([weights[v] for v in self.connected_verts[x]]) / len(self.connected_verts[x])
						average_weight = 0
						num_set = 0
						for v in self.connected_verts[x]:
							if not selected_only or obj.data.vertices[v].select:
								average_weight += weights[v]
								num_set += 1
						if num_set:
							average_weight /= num_set
						else:
							continue

					# Assign new weights with using laplace factor (or can I just call this a lerp? hmmm...)
					new_weights[i] = factor * average_weight + (1.0 - factor) * weights[i]

					# Grow or shrink weights.
					if blur_type == 1:
						new_weights[i] = min(new_weights[i], self.weights[i])
					elif blur_type == 2:
						new_weights[i] = max(new_weights[i], self.weights[i])

			weights = new_weights[:]

		# Update object weights
		for x, i in enumerate(self.vert_indexes):
			obj.data.vertices[i].groups[self.vert_group_indexes[x]].weight = weights[i]
		obj.data.update()

class WeightPaintBlurAll(bpy.types.Operator):
	bl_idname = "object.weightpaint_blur_all"
	bl_label = "Blur"
	bl_options = {'REGISTER', 'UNDO'}

	settings = bpy.props.PointerProperty(type = BlurSettingsCollection)

	active_index = None
	blur = None

	def draw(self, context):
		row = self.layout.row()
		row.prop(self.settings, "iterations", text = "Iter")
		row = self.layout.row()
		row.prop(self.settings, "selected_only")
		# self.layout.prop(self, "factor")
		box = self.layout.box()
		row = box.row()