diff --git a/role_expire.js b/role_expire.js
index 92ec894..c200105 100644
--- a/role_expire.js
+++ b/role_expire.js
@@ -6,31 +6,35 @@
  * Set of jQuery related routines.
  */
 
-$(document).ready(function() {
-  $('input.role-expire-role-expiry').parent().hide();
+(function ($) {
   
-  $('.role-expire-roles input.form-checkbox').each(function() {
-    var textfieldId = this.id.replace("roles", "role-expire");
-
-    // Move all expiry date fields under corresponding checkboxes
-    $(this).parent().after($('#'+textfieldId).parent());
-    
-    // Show all expiry date fields that have checkboxes checked
-    if ($(this).attr("checked")) {  
-      $('#'+textfieldId).parent().show();
-    }
-   
-  });
-  
-  $('.role-expire-roles input.form-checkbox').click(function() {
-    var textfieldId = this.id.replace("roles", "role-expire");
-    
-    // Toggle expiry date fields
-    if ($(this).attr("checked")) {
-      $('#'+textfieldId).parent().show(200);
-    } 
-    else {
-      $('#'+textfieldId).parent().hide(200);    
-    }
+  $(document).ready(function() {
+    $('input.role-expire-role-expiry').parent().hide();
+
+    $('#edit-roles input.form-checkbox').each(function() {
+      var textfieldId = this.id.replace("roles", "role-expire");
+
+      // Move all expiry date fields under corresponding checkboxes
+      $(this).parent().after($('#'+textfieldId).parent());
+
+      // Show all expiry date fields that have checkboxes checked
+      if ($(this).attr("checked")) {
+        $('#'+textfieldId).parent().show();
+      }
+
+    });
+
+    $('#edit-roles input.form-checkbox').click(function() {
+      var textfieldId = this.id.replace("roles", "role-expire");
+
+      // Toggle expiry date fields
+      if ($(this).attr("checked")) {
+        $('#'+textfieldId).parent().show(200);
+      }
+      else {
+        $('#'+textfieldId).parent().hide(200);
+      }
+    });
   });
-});
\ No newline at end of file
+
+})(jQuery);
\ No newline at end of file
diff --git a/role_expire.module b/role_expire.module
index 95414cc..f135739 100644
--- a/role_expire.module
+++ b/role_expire.module
@@ -39,7 +39,7 @@ function role_expire_get_all_user_records($uid) {
   $return = array();
   $result = db_query("SELECT rid, expiry_timestamp FROM {role_expire} WHERE uid = :uid", array(':uid' => $uid));
   foreach ($result as $row) {
-    $return[$row['rid']] = $row['expiry_timestamp'];
+    $return[$row->rid] = $row->expiry_timestamp;
   }
   return $return;
 }
@@ -234,7 +234,7 @@ function role_expire_form_user_register_alter(&$form, $form_state) {
  * Implements hook_form_FORM-ID_alter().
  */
 function role_expire_form_user_profile_form_alter(&$form, $form_state) {
-  $form['account']['roles']['#attributes'] = array('class' => 'role-expire-roles');
+  $form = array_merge_recursive($form, role_expire_add_expiration_input($form['#user']));
 }
 
 /**
@@ -248,7 +248,7 @@ function role_expire_form_user_admin_role_alter(&$form, $form_state) {
     '#size' => 8,
     '#default_value' => role_expire_get_default_duration($form['rid']['#value']),
     '#maxlength' => 5,
-    // TODO. WHY IS THIS CLASS BAD? '#attributes' => array('class' => 'role-expire-role-expiry'),
+    '#attributes' => array('class' => array('role-expire-role-expiry')),
     '#description' => t("Leave blank for no default length. Values should indicate the number of days the role should be active for.")
   );
   // Reposition the submit button and delete.
@@ -294,73 +294,37 @@ function role_expire_user_XXX($edit, $account) {
 }
 
 /**
- * Implements hook_user_presave().
+ * Implements hook_user_update().
  */
-function role_expire_user_presave(&$edit, $account, $category) {
-  if ($category == 'account' && (user_access('administer role expire') || user_access('administer users'))) {
-    if (isset($edit['roles'])) {
-      $time = REQUEST_TIME;
-      foreach ($edit['roles'] as $rid => $role) {
-        if (is_array($edit) && array_key_exists('role_expire_' . $role, $edit) && $edit['role_expire_' . $role] != '') {
-          $expiry_time = strtotime($edit['role_expire_' . $role]);
-          if (!$expiry_time) {
-            form_set_error('role_expire_' . $role, t("Role expiry time is not in correct format."));
-          }
-          if ($expiry_time <= $time) {
-            form_set_error('role_expire_' . $role, t("Role expiry time must be in the future."));
-          }
-        }
-      }
-    }
-  }
+function role_expire_user_update(&$edit, $account, $category) {
   if ($category == 'account' && (user_access('administer role expire') || user_access('administer users'))) {
+
     // Add roles expiry information for the user role.
-    // We go over all existing roles, because user might have disabled a role.
-    foreach (_role_expire_get_role() as $rid => $role) {
-      $role = '';
-      if (array_key_exists('role_expire_' . $rid, $edit)) {
-        $role = $edit['role_expire_' . $rid];
-      }
-      if ($role != '') {
-        $expiry_timestamp = strtotime($role);
-        role_expire_write_record($account->uid, $rid, $expiry_timestamp);
-      }
-      else {
-        if ($category == NULL || $category == 'account') {
-          // Handle inserts correctly
-          if ($op == 'insert') {
-            // Check for Role
-            if (array_key_exists($rid, $edit['roles']) && $edit['roles'][$rid] == $rid) {
-              role_expire_process_default_role_duration_for_user($rid, $account->uid);
-            }
-          }
-          else {
-            // User input is blank, so delete record.
-            role_expire_delete_record($account->uid, $rid);
-          }
+    foreach (array_keys($edit) as $name) {
+      if (strpos($name, 'role_expire_') === 0) {
+        $value = $edit[$name];
+        $rid = substr($name, strlen('role_expire_'));
+        if ($value != '' && array_key_exists($rid, $edit['roles'])) {
+          $expiry_timestamp = strtotime($value);
+          role_expire_write_record($account->uid, $rid, $expiry_timestamp);
+        }
+        else {
+          role_expire_delete_record($account->uid, $rid);
         }
       }
-      $edit['role_expire_' . $rid] = NULL;
     }
-  }
-}
-
-/**
- * Implements hook_user_update().
- */
-function role_expire_user_update(&$edit, $account, $category) {
-  if ($category == 'account' && (user_access('administer role expire') || user_access('administer users'))) {
-    // This adds any new roles that have been given to the user.
-    $new_roles = array_diff($account->roles, _role_static_user_roles($account->uid));
+
+    // This adds default expiration to any new roles that have been given to the user.
+    $new_roles = array_diff(array_keys($edit['roles']), array_keys($edit['original']->roles));
     // We have the new roles, loop over them and see whether we need to assign expiry to them.
-    foreach ($new_roles as $role_id => $role) {
+    foreach ($new_roles as $role_id) {
       role_expire_process_default_role_duration_for_user($role_id, $account->uid);
     }
 
-    // This removes any new roles that have been given to the user.
-    $del_roles = array_diff(_role_static_user_roles($account->uid), $account->roles);
+    // This removes expiration for roles that have been removed from the user.
+    $del_roles = array_diff(array_keys($edit['original']->roles), array_keys($edit['roles']));
     // We have the deleted roles, loop over them and see whether we need to assign expiry to them.
-    foreach ($del_roles as $role_id => $role) {
+    foreach ($del_roles as $role_id) {
       role_expire_delete_record($account->uid, $role_id);
     }
   }
@@ -372,34 +336,14 @@ function role_expire_user_update(&$edit, $account, $category) {
  * Implements hook_user_insert().
  */
 function role_expire_user_insert(&$edit, $account, $category) {
+
   if ($category == 'account' && (user_access('administer role expire') || user_access('administer users'))) {
-    // Add roles expiry information for the user role.
-    // We go over all existing roles, because user might have disabled a role.
-    foreach (_role_expire_get_role() as $rid => $role) {
-      $role = '';
-      if (array_key_exists('role_expire_' . $rid, $edit)) {
-        $role = $edit['role_expire_' . $rid];
-      }
-      if ($role != '') {
-        $expiry_timestamp = strtotime($role);
-        role_expire_write_record($account->uid, $rid, $expiry_timestamp);
-      }
-      else {
-        if ($category == NULL || $category == 'account') {
-          // Handle inserts correctly
-          if ($op == 'insert') {
-            // Check for Role
-            if (array_key_exists($rid, $edit['roles']) && $edit['roles'][$rid] == $rid) {
-              role_expire_process_default_role_duration_for_user($rid, $account->uid);
-            }
-          }
-          else {
-            // User input is blank, so delete record.
-            role_expire_delete_record($account->uid, $rid);
-          }
-        }
-      }
-      $edit['role_expire_' . $rid] = NULL;
+
+    // This adds default expiration to any new roles that have been given to the user.
+    $new_roles = array_keys($edit['roles']);
+    // We have the new roles, loop over them and see whether we need to assign expiry to them.
+    foreach ($new_roles as $role_id) {
+      role_expire_process_default_role_duration_for_user($role_id, $account->uid);
     }
   }
 }
@@ -449,8 +393,8 @@ function role_expire_user_view($account, $view_mode) {
       $account->content['summary']['role_expire'] = array(
           '#type' => 'user_profile_item',
           '#title' => t('Role expiration'),
-          '#value' => theme('item_list', array('items' => $roles)),
-          '#attributes' => array('class' => 'role-expiry-roles'),
+          '#markup' => theme('item_list', array('items' => $roles)),
+          '#attributes' => array('class' => array('role-expiry-roles')),
         );
     }
   }
@@ -501,7 +445,7 @@ function role_expire_add_expiration_input($account = NULL) {
   $form = array();
   if (user_access('administer users') || user_access('administer role expire')) {
     drupal_add_js(drupal_get_path('module', 'role_expire') . '/role_expire.js');
-    $form['roles']['#attributes'] = array('class' => 'role-expire-roles');
+    $form['roles']['#attributes'] = array('class' => array('role-expire-roles'));
 
     foreach (_role_expire_get_role() as $rid => $role) {
       if (is_object($account) and array_key_exists('uid', $account)) {
@@ -510,14 +454,17 @@ function role_expire_add_expiration_input($account = NULL) {
       else {
         $expiry_timestamp = '';
       }
+
       $form['role_expire_' . $rid] = array(
         '#title' => t("%role role expiration time", array('%role' => drupal_ucfirst($role))),
         '#type' => 'textfield',
         '#default_value' => !empty($expiry_timestamp) ? date("d-m-Y G:i:s", $expiry_timestamp) : '',
-        '#attributes' => array('class' => 'role-expire-role-expiry'),
+        '#attributes' => array('class' => array('role-expire-role-expiry')),
         '#description' => t("Leave blank for default role expiry (never, or the duration you have set for the role), enter date and time in format: <em>yyyy-mm-dd hh:mm:ss</em> or use relative time i.e. 1 day, 2 months, 1 year, 3 years.")
       );
     }
+
+    $form['#validate'][] = '_role_expire_validate_role_expires';
   }
   return $form;
 }
@@ -556,6 +503,29 @@ function _role_expire_get_role() {
 }
 
 /**
+ * Form validation handler for the role expiration on the user_profile_form().
+ *
+ * @see user_profile_form()
+ */
+function _role_expire_validate_role_expires(&$form, &$form_state) {
+
+  $time = REQUEST_TIME;
+
+  foreach ($form_state['values'] as $name => $value) {
+    if (strpos($name, 'role_expire_') === 0 && trim($value) != '')
+    {
+      $expiry_time = strtotime($value);
+      if (!$expiry_time) {
+        form_set_error($name, t("Role expiry time is not in correct format."));
+      }
+      if ($expiry_time <= $time) {
+        form_set_error($name, t("Role expiry time must be in the future."));
+      }
+    }
+  }
+}
+
+/**
  * @todo Please document this function.
  * @see http://drupal.org/node/1354
  */