diff --git a/autosave.module b/autosave.module
index 1693eef..252eef5 100644
--- a/autosave.module
+++ b/autosave.module
@@ -30,13 +30,14 @@ function autosave_menu() {
     'access callback' => 'autosave_save_access',
     'type'            => MENU_CALLBACK,
   );
+
   $items['autosave/restore'] = array(
     'title'           => 'Autosave form restore',
     'page callback'   => 'autosave_restore',
     'access callback' => 'autosave_restore_access',
     'access arguments' => array(2, 3),
     'type'            => MENU_CALLBACK,
-    'delivery callback' => 'ajax_deliver',
+    'delivery callback' => 'ajax_deliver'
   );
 
   $items['admin/config/autosave'] = array(
@@ -79,7 +80,7 @@ function autosave_admin_settings($form, &$form_state) {
 /**
  * Implements hook_form_alter() for node_type_form().
  */
-function autosave_form_node_type_form_alter(&$form, $form_state) {
+function autosave_form_node_type_form_alter(&$form, &$form_state) {
   $form['workflow']['autosave'] = array(
     '#type' => 'checkbox',
     '#title' => t('Enable Autosave to add/edit forms for this node type'),
@@ -95,8 +96,15 @@ function autosave_form_alter(&$form, &$form_state, $form_id) {
   global $user;
   $path = request_path();
 
-  if (stristr($form_id, '_node_form') && arg(0) != 'admin') {
+	// Add custom js
+	if (stristr($form_id, '_node_form') && arg(0) != 'admin') {
+  	if ((variable_get('autosave_' . $form['type']['#value'], 0)) && empty($_POST['autosave_form_path'])) {
+    	drupal_add_js(AUTOSAVE_PATH . '/autosave_custom_elements.js');
+		}
+	}
 
+  // CUSTOM: added !$form_state['rebuild'] - experimental
+  if (stristr($form_id, '_node_form') && arg(0) != 'admin' && !$form_state['rebuild']) {
     // Check if this content_type has the autosave function enabled and make sure it's a node edit or add form.
     if ((variable_get('autosave_' . $form['type']['#value'], 0)) && empty($_POST['autosave_form_path'])) {
       drupal_add_library('system', 'drupal.ajax');
@@ -131,6 +139,18 @@ function autosave_form_alter(&$form, &$form_state, $form_id) {
 
       drupal_add_js($settings, array('type' => 'setting', 'scope' => JS_DEFAULT));
     }
+  } elseif ($form_state['rebuild']) {
+		
+    // CUSTOM: hack for hierarchical select - reset field_categories value, because value lost after form rebuild
+		if(!empty($form_state['storage']['hs_custom'])) {
+			foreach($form_state['storage']['hs_custom'] as $key => $value) {
+				$form_state['storage']['hs'][] = $value[0];
+				$form_state['input'][$key] = $value[0]['dropbox_lineages_selections'];
+			}
+		}
+
+    $settings['autosave']['savedTimestamp'] = 0;
+    drupal_add_js($settings, array('type' => 'setting', 'scope' => JS_DEFAULT));
   }
 }
 
@@ -145,9 +165,12 @@ function autosave_form_alter(&$form, &$form_state, $form_id) {
  *   The timestamp at which the autosaved form was saved.  This is used to
  *   differentiate between different people mucking with the same form.
  */
+
 function autosave_restore($formid, $timestamp) {
   global $user;
 
+  $commands = Array();
+  
   // Convert the form ID back to the PHP version.  I do hate that duality...
   $form_id = str_replace("-", "_", $formid);
 
@@ -163,11 +186,37 @@ function autosave_restore($formid, $timestamp) {
       require_once DRUPAL_ROOT . '/' . $menu_item['include_file'];
     }
     $form_state['input'] = unserialize($record->serialized);
-    $form_state['build_info']['args'] = $menu_item['page_arguments'];
+		// CUSTOM : restoring an existin node form
+		if(is_object($menu_item['page_arguments'][0]) && isset($menu_item['page_arguments'][0]->nid)) {
+    	$form_state['build_info']['args'] = $menu_item['page_arguments'];
+		// CUSTOM : restoring a node/add form
+		} else {
+			$node = (object) array('uid' => $user->uid, 'name' => (isset($user->name) ? $user->name : ''), 'type' => $menu_item['page_arguments'][0], 'language' => LANGUAGE_NONE);
+			$form_state['build_info']['args'][] = $node;
+		}
+			
     // Disable the "this form has already been submitted" nonsense by making
     // Drupal think the form is being rebuilt as part of a multi-step form.
     $form_state['rebuild'] = TRUE;
-
+		
+		// CUSTOM: ensure all hs-type fields are saved in our custom storage
+		$fields = field_info_instances('node', $form_state['build_info']['args'][0]->type);
+		$check_fields = array();
+		if(!empty($fields)) {
+			foreach($fields as $field_name => $value) {
+				if($value['widget']['type'] == 'taxonomy_hs') {
+					$check_fields[] = $field_name;
+				}
+			}
+		}
+		
+		if(!empty($check_fields)) {
+			foreach($check_fields as $field) {
+				// CUSTOM: added extra storage for field_categories, couse 'input' storage gets "stripped" after form rebuild
+				$form_state['storage']['hs_custom'][$field] = array(array('dropbox_lineages_selections' => $form_state['input'][$field]));
+			}
+		}
+    
     $form = drupal_build_form($form_id, $form_state);
 
     // Because the form will by default submit back to this URL, we need to
@@ -175,7 +224,7 @@ function autosave_restore($formid, $timestamp) {
     // originally.
     $form['#action'] = url($record->path);
     $form['autosave_form_path']['#value'] = $form['#action'];
-
+    
     // We don't want to change the HTML ID of the form, because we're replacing
     // it in-place.  Drupal wants to give this a suffix for some reason.
     $form['#id'] = $formid;
@@ -183,9 +232,9 @@ function autosave_restore($formid, $timestamp) {
     $commands[] = ajax_command_replace('#' . $formid, drupal_render($form));
 
     return array('#type' => 'ajax', '#commands' => $commands);
-
-    return $commands;
   }
+  
+  return $commands;
 }
 
 /**
@@ -220,6 +269,7 @@ function autosave_save_access() {
  * @return boolean
  *   True if the user should have restore access to this form, false otherwise.
  */
+
 function autosave_restore_access($formid, $timestamp) {
   $record = autosave_get_autosaved_form($formid, $timestamp, $GLOBALS['user']->uid);
 
@@ -240,8 +290,7 @@ function autosave_save() {
   // Not all variables need to be serialized.
   //    - for Drupal 6 version need to remove op and form_build_id
   unset($_POST['autosave_form_path'], $_POST['op'], $_POST['form_build_id']);
-  $serialized = serialize($_POST);
-
+  
   // check if node has just been saved - if it has then it's because AS ajax fired off as user was submitting
   // if it had just been submitted - no need to AS now
   //    - easy to figure out if we are submitting an edit to existing node
@@ -250,15 +299,55 @@ function autosave_save() {
   // update case
   if (is_numeric($path_args[1])) {
     $submitted = node_load($path_args[1]);
+		$type = $submitted->type;
   }
   else {
     // add case
+		$type = str_replace("-", "_", $path_args[2]);
     $submitted->changed = db_query("SELECT created FROM {node} WHERE uid = :uid and type = :type ORDER BY created DESC LIMIT 1", array(
       ':uid' => $user->uid,
-      ':type' => str_replace("-", "_", $path_args[2]))
+      ':type' => $type)
     )->fetchField();
   }
 
+	// CUSTOM: Retrieve all field instances with the hierarchical_select widget attached to this node type
+	$fields = field_info_instances('node', $type);
+	$check_fields = array();
+	if(!empty($fields)) {
+		foreach($fields as $field_name => $value) {
+			if($value['widget']['type'] == 'taxonomy_hs') {
+				$check_fields[] = $field_name;
+			}
+		}
+	}
+	
+  // CUSTOM: hack - set HS values manually
+  $copyofPost = $_POST;
+	
+	// CUSTOM: handle storage of HS values
+	if(!empty($check_fields)) {
+		foreach($check_fields as $field) {
+			// TODO: a cleaner way to determine field language?
+			$lang = isset($submitted->{$field}['lv']) ? 'lv' : 'und';
+			
+			// categories already saved, just passing them along
+		  if (!empty($submitted->{$field}[$lang])) {
+		    foreach ($submitted->{$field}[$lang] as $key => $catObj) {
+		      $copyofPost[$field][$lang][$key] = $catObj['tid'];
+		    }
+		  }
+		
+			// categories added to dropbox, but not saved yet
+		  if (!empty($copyofPost[$field.'_extra_categories'])) {
+				foreach($copyofPost[$field.'_extra_categories'] as $key => $value) {
+					$copyofPost[$field][$lang][] = $value;
+				}
+		  }
+		}
+	}
+	
+  $serialized = serialize($copyofPost);
+  
   if (!$submitted || (REQUEST_TIME - $submitted->changed) > 10) {
     // Currently, each user can have only one autosave form at a particular path.
     db_merge('autosaved_forms')->key(array(
@@ -292,7 +381,6 @@ function autosave_save() {
  *   An array containing the serialized values of the autosaved form and the timestamp of when the form was autosaved.
  */
 function autosave_get_autosaved_form($form_id, $timestamp, $uid) {
-
   static $forms = array();
 
   if (empty($forms[$form_id][$timestamp])) {
@@ -304,7 +392,6 @@ function autosave_get_autosaved_form($form_id, $timestamp, $uid) {
     ))->fetchObject();
   }
 
-
   return $forms[$form_id][$timestamp];
 }